@@ -102,13 +102,16 @@ func (s AccessEntryService) RecordDecision(
|
||||
entry.DecisionNote = req.DecisionNote
|
||||
entry.DecidedBy = req.DecidedByID
|
||||
entry.DecidedAt = &now
|
||||
|
||||
entry.UpdatedAt = now
|
||||
if entry.Flags == nil {
|
||||
entry.Flags = []coredata.AccessEntryFlag{}
|
||||
}
|
||||
|
||||
if entry.FlagReasons == nil {
|
||||
entry.FlagReasons = []string{}
|
||||
}
|
||||
|
||||
if req.Decision == coredata.AccessEntryDecisionRevoke || req.Decision == coredata.AccessEntryDecisionEscalate {
|
||||
if len(entry.Flags) == 0 {
|
||||
entry.Flags = []coredata.AccessEntryFlag{coredata.AccessEntryFlagExcessive}
|
||||
@@ -156,6 +159,7 @@ func (s AccessEntryService) RecordDecisions(
|
||||
if d.Decision == coredata.AccessEntryDecisionPending {
|
||||
return nil, fmt.Errorf("cannot bulk decide access entries: invalid decision %q", d.Decision)
|
||||
}
|
||||
|
||||
if d.Decision != coredata.AccessEntryDecisionApproved {
|
||||
if d.DecisionNote == nil || strings.TrimSpace(*d.DecisionNote) == "" {
|
||||
return nil, fmt.Errorf(
|
||||
@@ -189,9 +193,11 @@ func (s AccessEntryService) RecordDecisions(
|
||||
if err := campaign.LoadByID(ctx, conn, s.scope, entry.AccessReviewCampaignID); err != nil {
|
||||
return fmt.Errorf("cannot load campaign: %w", err)
|
||||
}
|
||||
|
||||
if campaign.Status != coredata.AccessReviewCampaignStatusPendingActions {
|
||||
return fmt.Errorf("cannot decide access entry: campaign status is %s, expected PENDING_ACTIONS", campaign.Status)
|
||||
}
|
||||
|
||||
verifiedCampaigns[entry.AccessReviewCampaignID] = true
|
||||
}
|
||||
|
||||
@@ -200,13 +206,16 @@ func (s AccessEntryService) RecordDecisions(
|
||||
entry.DecisionNote = d.DecisionNote
|
||||
entry.DecidedBy = d.DecidedByID
|
||||
entry.DecidedAt = &now
|
||||
|
||||
entry.UpdatedAt = now
|
||||
if entry.Flags == nil {
|
||||
entry.Flags = []coredata.AccessEntryFlag{}
|
||||
}
|
||||
|
||||
if entry.FlagReasons == nil {
|
||||
entry.FlagReasons = []string{}
|
||||
}
|
||||
|
||||
if d.Decision == coredata.AccessEntryDecisionRevoke || d.Decision == coredata.AccessEntryDecisionEscalate {
|
||||
if len(entry.Flags) == 0 {
|
||||
entry.Flags = []coredata.AccessEntryFlag{coredata.AccessEntryFlagExcessive}
|
||||
@@ -245,6 +254,7 @@ func (s AccessEntryService) RecordDecisions(
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot reload access entry %s: %w", id, err)
|
||||
}
|
||||
|
||||
entries[i] = entry
|
||||
}
|
||||
|
||||
@@ -274,14 +284,17 @@ func (s AccessEntryService) FlagEntry(
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
|
||||
entry.Flags = req.Flags
|
||||
if entry.Flags == nil {
|
||||
entry.Flags = []coredata.AccessEntryFlag{}
|
||||
}
|
||||
|
||||
entry.FlagReasons = req.FlagReasons
|
||||
if entry.FlagReasons == nil {
|
||||
entry.FlagReasons = []string{}
|
||||
}
|
||||
|
||||
entry.UpdatedAt = now
|
||||
|
||||
return entry.UpdateFlags(ctx, conn, s.scope)
|
||||
@@ -348,10 +361,12 @@ func (s AccessEntryService) CountForCampaignID(
|
||||
ctx,
|
||||
func(ctx context.Context, conn pg.Querier) (err error) {
|
||||
entries := coredata.AccessEntries{}
|
||||
|
||||
count, err = entries.CountByCampaignID(ctx, conn, s.scope, campaignID, filter)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot count access entries by campaign: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
)
|
||||
@@ -374,10 +389,12 @@ func (s AccessEntryService) CountForCampaignIDAndSourceID(
|
||||
ctx,
|
||||
func(ctx context.Context, conn pg.Querier) (err error) {
|
||||
entries := coredata.AccessEntries{}
|
||||
|
||||
count, err = entries.CountByCampaignIDAndSourceID(ctx, conn, s.scope, campaignID, sourceID, filter)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot count access entries by campaign and source: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
)
|
||||
@@ -398,10 +415,12 @@ func (s AccessEntryService) CountPendingForCampaignID(
|
||||
ctx,
|
||||
func(ctx context.Context, conn pg.Querier) (err error) {
|
||||
entries := coredata.AccessEntries{}
|
||||
|
||||
count, err = entries.CountPendingByCampaignID(ctx, conn, s.scope, campaignID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot count pending access entries: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
)
|
||||
|
||||
@@ -188,6 +188,7 @@ func (s AccessSourceService) Update(
|
||||
return fmt.Errorf("cannot load connector: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
source.ConnectorID = *req.ConnectorID
|
||||
}
|
||||
|
||||
@@ -256,6 +257,7 @@ func (s AccessSourceService) CountForOrganizationID(
|
||||
func(ctx context.Context, conn pg.Querier) (err error) {
|
||||
sources := coredata.AccessSources{}
|
||||
count, err = sources.CountByOrganizationID(ctx, conn, s.scope, organizationID)
|
||||
|
||||
return err
|
||||
},
|
||||
)
|
||||
@@ -300,6 +302,7 @@ func (s AccessSourceService) ConnectorHTTPClient(
|
||||
if err := dbConnector.LoadByID(ctx, conn, s.scope, connectorID, s.encryptionKey); err != nil {
|
||||
return fmt.Errorf("cannot load connector: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
)
|
||||
@@ -308,16 +311,19 @@ func (s AccessSourceService) ConnectorHTTPClient(
|
||||
}
|
||||
|
||||
var tokenBefore string
|
||||
|
||||
oauth2Conn, isOAuth2 := dbConnector.Connection.(*connector.OAuth2Connection)
|
||||
if isOAuth2 {
|
||||
tokenBefore = oauth2Conn.AccessToken
|
||||
}
|
||||
|
||||
var httpClient *http.Client
|
||||
|
||||
if isOAuth2 && s.connectorRegistry != nil {
|
||||
refreshCfg := s.connectorRegistry.GetOAuth2RefreshConfig(string(dbConnector.Provider))
|
||||
if refreshCfg != nil {
|
||||
var err error
|
||||
|
||||
httpClient, err = oauth2Conn.RefreshableClient(ctx, *refreshCfg)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("cannot create refreshable HTTP client: %w", err)
|
||||
@@ -327,6 +333,7 @@ func (s AccessSourceService) ConnectorHTTPClient(
|
||||
|
||||
if httpClient == nil {
|
||||
var err error
|
||||
|
||||
httpClient, err = dbConnector.Connection.Client(ctx)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("cannot create HTTP client: %w", err)
|
||||
@@ -336,6 +343,7 @@ func (s AccessSourceService) ConnectorHTTPClient(
|
||||
// Persist refreshed token if it changed.
|
||||
if isOAuth2 && oauth2Conn.AccessToken != tokenBefore {
|
||||
dbConnector.UpdatedAt = time.Now()
|
||||
|
||||
if err := s.pg.WithTx(
|
||||
ctx,
|
||||
func(ctx context.Context, tx pg.Tx) error {
|
||||
|
||||
@@ -105,6 +105,7 @@ func (s *CampaignService) Get(
|
||||
if err := campaign.LoadByID(ctx, conn, s.scope, campaignID); err != nil {
|
||||
return fmt.Errorf("cannot load campaign: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
)
|
||||
@@ -192,6 +193,7 @@ func (s *CampaignService) Delete(
|
||||
if err := campaign.Delete(ctx, conn, s.scope); err != nil {
|
||||
return fmt.Errorf("cannot delete campaign: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
)
|
||||
@@ -359,6 +361,7 @@ func (s *CampaignService) Close(
|
||||
}
|
||||
|
||||
entries := coredata.AccessEntries{}
|
||||
|
||||
pendingCount, err := entries.CountPendingByCampaignID(ctx, conn, s.scope, campaignID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot count pending entries: %w", err)
|
||||
@@ -392,6 +395,7 @@ func lockCampaignForUpdate(ctx context.Context, tx pg.Tx, scope coredata.Scoper,
|
||||
if err := c.LockForUpdate(ctx, tx, scope); err != nil {
|
||||
return fmt.Errorf("cannot lock campaign for update: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -402,6 +406,7 @@ func (s *CampaignService) enqueueSourceFetches(
|
||||
sources coredata.AccessSources,
|
||||
) error {
|
||||
now := time.Now()
|
||||
|
||||
for _, source := range sources {
|
||||
fetch := &coredata.AccessReviewCampaignSourceFetch{
|
||||
AccessReviewCampaignID: campaignID,
|
||||
@@ -469,6 +474,7 @@ func (s *CampaignService) ListForOrganizationID(
|
||||
if err := campaigns.LoadByOrganizationID(ctx, conn, s.scope, organizationID, cursor); err != nil {
|
||||
return fmt.Errorf("cannot load campaigns by organization: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
)
|
||||
@@ -491,6 +497,7 @@ func (s *CampaignService) ListSourceFetches(
|
||||
if err := fetches.LoadByCampaignID(ctx, conn, s.scope, campaignID); err != nil {
|
||||
return fmt.Errorf("cannot load source fetches by campaign: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
)
|
||||
@@ -511,6 +518,7 @@ func (s *CampaignService) CountForOrganizationID(
|
||||
ctx,
|
||||
func(ctx context.Context, conn pg.Querier) (err error) {
|
||||
campaigns := coredata.AccessReviewCampaigns{}
|
||||
|
||||
count, err = campaigns.CountByOrganizationID(ctx, conn, s.scope, organizationID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot count campaigns by organization: %w", err)
|
||||
|
||||
@@ -98,6 +98,7 @@ func (d *AsanaDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error)
|
||||
if page.NextPage == nil || page.NextPage.URI == "" {
|
||||
return records, nil
|
||||
}
|
||||
|
||||
next = page.NextPage.URI
|
||||
}
|
||||
|
||||
@@ -109,12 +110,14 @@ func (d *AsanaDriver) queryUsers(ctx context.Context, endpoint string) (*asanaUs
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create asana users request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute asana users request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
|
||||
@@ -110,12 +110,14 @@ func (d *BitbucketDriver) queryMembers(ctx context.Context, endpoint string) (*b
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create bitbucket members request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute bitbucket members request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
|
||||
@@ -84,6 +84,7 @@ func (d *BrexDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error)
|
||||
if resp.NextCursor == "" {
|
||||
return records, nil
|
||||
}
|
||||
|
||||
nextCursor := resp.NextCursor
|
||||
cursor = &nextCursor
|
||||
}
|
||||
@@ -110,6 +111,7 @@ func (d *BrexDriver) queryUsers(ctx context.Context, cursor *string) (*brexUsers
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute brex users request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
|
||||
@@ -82,6 +82,7 @@ func TestCassettesUseSyntheticEmails(t *testing.T) {
|
||||
if seen[email] {
|
||||
continue
|
||||
}
|
||||
|
||||
seen[email] = true
|
||||
|
||||
domain := email[strings.IndexByte(email, '@')+1:]
|
||||
@@ -90,6 +91,7 @@ func TestCassettesUseSyntheticEmails(t *testing.T) {
|
||||
}
|
||||
|
||||
ok := false
|
||||
|
||||
for _, suffix := range allowedDomainSuffixes {
|
||||
if strings.HasSuffix("."+domain, suffix) || domain == strings.TrimPrefix(suffix, ".") {
|
||||
ok = true
|
||||
|
||||
@@ -77,12 +77,14 @@ func (d *ClickUpDriver) ListAccounts(ctx context.Context) ([]AccountRecord, erro
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create clickup team request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute clickup team request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -159,5 +161,6 @@ func parseClickUpTime(raw string) (time.Time, error) {
|
||||
if err != nil {
|
||||
return time.Time{}, fmt.Errorf("cannot parse clickup time %q: %w", raw, err)
|
||||
}
|
||||
|
||||
return time.UnixMilli(ms).UTC(), nil
|
||||
}
|
||||
|
||||
@@ -131,6 +131,7 @@ func (d *CloudflareDriver) queryAccounts(ctx context.Context, page int) (*cloudf
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute cloudflare accounts request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
@@ -168,6 +169,7 @@ func (d *CloudflareDriver) queryAllMembers(ctx context.Context, accountID string
|
||||
}
|
||||
|
||||
isAdmin := false
|
||||
|
||||
for _, r := range m.Roles {
|
||||
if r.Name == "Super Administrator - All Privileges" || r.Name == "Administrator" {
|
||||
isAdmin = true
|
||||
@@ -224,6 +226,7 @@ func (d *CloudflareDriver) queryMembers(ctx context.Context, accountID string, p
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute cloudflare members request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
|
||||
@@ -51,6 +51,7 @@ func (d *CSVDriver) ListAccounts(_ context.Context) ([]AccountRecord, error) {
|
||||
for i, col := range header {
|
||||
colIndex[strings.TrimSpace(strings.ToLower(col))] = i
|
||||
}
|
||||
|
||||
if _, ok := colIndex["email"]; !ok {
|
||||
return nil, fmt.Errorf("cannot parse CSV: missing required column email")
|
||||
}
|
||||
@@ -62,6 +63,7 @@ func (d *CSVDriver) ListAccounts(_ context.Context) ([]AccountRecord, error) {
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read CSV row: %w", err)
|
||||
}
|
||||
@@ -75,24 +77,31 @@ func (d *CSVDriver) ListAccounts(_ context.Context) ([]AccountRecord, error) {
|
||||
if idx, ok := colIndex["email"]; ok && idx < len(row) {
|
||||
record.Email = strings.TrimSpace(row[idx])
|
||||
}
|
||||
|
||||
if idx, ok := colIndex["full_name"]; ok && idx < len(row) {
|
||||
record.FullName = strings.TrimSpace(row[idx])
|
||||
}
|
||||
|
||||
if idx, ok := colIndex["role"]; ok && idx < len(row) {
|
||||
record.Role = strings.TrimSpace(row[idx])
|
||||
}
|
||||
|
||||
if idx, ok := colIndex["job_title"]; ok && idx < len(row) {
|
||||
record.JobTitle = strings.TrimSpace(row[idx])
|
||||
}
|
||||
|
||||
if idx, ok := colIndex["is_admin"]; ok && idx < len(row) {
|
||||
record.IsAdmin = strings.TrimSpace(strings.ToLower(row[idx])) == "true"
|
||||
}
|
||||
|
||||
if idx, ok := colIndex["active"]; ok && idx < len(row) {
|
||||
record.Active = new(strings.TrimSpace(strings.ToLower(row[idx])) == "true")
|
||||
}
|
||||
|
||||
if idx, ok := colIndex["external_id"]; ok && idx < len(row) {
|
||||
record.ExternalID = strings.TrimSpace(row[idx])
|
||||
}
|
||||
|
||||
if idx, ok := colIndex["account_type"]; ok && idx < len(row) {
|
||||
if strings.TrimSpace(strings.ToUpper(row[idx])) == "SERVICE_ACCOUNT" {
|
||||
record.AccountType = coredata.AccessEntryAccountTypeServiceAccount
|
||||
|
||||
@@ -24,6 +24,7 @@ func TestCSVDriverRequiresEmailHeader(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
driver := NewCSVDriver(strings.NewReader("full_name,role\nJane Doe,Admin\n"))
|
||||
|
||||
_, err := driver.ListAccounts(context.Background())
|
||||
if err == nil {
|
||||
t.Fatalf("expected error when email header is missing")
|
||||
@@ -36,16 +37,20 @@ func TestCSVDriverParsesRequiredAndOptionalColumns(t *testing.T) {
|
||||
driver := NewCSVDriver(strings.NewReader(
|
||||
"email,full_name,role,external_id\njane@example.com,Jane Doe,Admin,42\n",
|
||||
))
|
||||
|
||||
records, err := driver.ListAccounts(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if len(records) != 1 {
|
||||
t.Fatalf("expected 1 record, got %d", len(records))
|
||||
}
|
||||
|
||||
if records[0].Email != "jane@example.com" {
|
||||
t.Fatalf("unexpected email: %s", records[0].Email)
|
||||
}
|
||||
|
||||
if records[0].ExternalID != "42" {
|
||||
t.Fatalf("unexpected external id: %s", records[0].ExternalID)
|
||||
}
|
||||
|
||||
@@ -79,6 +79,7 @@ func (d *DocuSignDriver) ListAccounts(ctx context.Context) ([]AccountRecord, err
|
||||
}
|
||||
|
||||
var records []AccountRecord
|
||||
|
||||
startPosition := 0
|
||||
|
||||
for range maxPaginationPages {
|
||||
@@ -143,12 +144,14 @@ func (d *DocuSignDriver) discoverAccount(ctx context.Context) (accountID string,
|
||||
if err != nil {
|
||||
return "", "", fmt.Errorf("cannot create docusign userinfo request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", "", fmt.Errorf("cannot execute docusign userinfo request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
@@ -183,12 +186,14 @@ func (d *DocuSignDriver) queryUsers(ctx context.Context, baseURI string, account
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create docusign users request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute docusign users request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
|
||||
@@ -85,6 +85,7 @@ func (d *GitHubDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error
|
||||
d.logger.WarnCtx(ctx, "cannot fetch github membership, skipping member",
|
||||
log.Error(err),
|
||||
)
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -93,6 +94,7 @@ func (d *GitHubDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error
|
||||
d.logger.WarnCtx(ctx, "cannot fetch github user profile, skipping member",
|
||||
log.Error(err),
|
||||
)
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -107,6 +109,7 @@ func (d *GitHubDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error
|
||||
}
|
||||
|
||||
mfaStatus := coredata.MFAStatusUnknown
|
||||
|
||||
if no2FASet != nil {
|
||||
if no2FASet[m.Login] {
|
||||
mfaStatus = coredata.MFAStatusDisabled
|
||||
@@ -158,6 +161,7 @@ func (d *GitHubDriver) fetchAllMembers(ctx context.Context) ([]githubMember, err
|
||||
if nextURL == "" {
|
||||
return members, nil
|
||||
}
|
||||
|
||||
url = nextURL
|
||||
}
|
||||
|
||||
@@ -176,6 +180,7 @@ func (d *GitHubDriver) fetchMembersPage(ctx context.Context, url string) ([]gith
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("cannot execute github members request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
@@ -215,6 +220,7 @@ func (d *GitHubDriver) fetchAll2FADisabledLogins(ctx context.Context) (map[strin
|
||||
if nextURL == "" {
|
||||
return set, nil
|
||||
}
|
||||
|
||||
url = nextURL
|
||||
}
|
||||
|
||||
@@ -239,6 +245,7 @@ func (d *GitHubDriver) fetchMembership(ctx context.Context, login string) (*gith
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute github membership request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
@@ -269,6 +276,7 @@ func (d *GitHubDriver) fetchUserProfile(ctx context.Context, login string) (*git
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute github user profile request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
|
||||
@@ -117,12 +117,14 @@ func (d *GitLabDriver) queryMembers(ctx context.Context, endpoint string) ([]git
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("cannot create gitlab members request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("cannot execute gitlab members request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
|
||||
@@ -59,6 +59,7 @@ func (rt *retryRoundTripper) RoundTrip(req *http.Request) (*http.Response, error
|
||||
}
|
||||
|
||||
var lastResp *http.Response
|
||||
|
||||
for attempt := range rt.maxRetries {
|
||||
resp, err := transport.RoundTrip(req)
|
||||
if err != nil {
|
||||
@@ -94,6 +95,7 @@ func (d *GoogleWorkspaceDriver) ListAccounts(ctx context.Context) ([]AccountReco
|
||||
}
|
||||
|
||||
var records []AccountRecord
|
||||
|
||||
pageToken := ""
|
||||
|
||||
for range maxPaginationPages {
|
||||
|
||||
@@ -124,6 +124,7 @@ func (d *HerokuDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error
|
||||
if nextRange == "" {
|
||||
return records, nil
|
||||
}
|
||||
|
||||
rangeHeader = nextRange
|
||||
}
|
||||
|
||||
@@ -135,7 +136,9 @@ func (d *HerokuDriver) queryMembers(ctx context.Context, endpoint, rangeHeader s
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("cannot create heroku members request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/vnd.heroku+json; version=3")
|
||||
|
||||
if rangeHeader != "" {
|
||||
req.Header.Set("Range", rangeHeader)
|
||||
}
|
||||
@@ -144,6 +147,7 @@ func (d *HerokuDriver) queryMembers(ctx context.Context, endpoint, rangeHeader s
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("cannot execute heroku members request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
// Heroku returns 206 Partial Content for ranged responses with more
|
||||
|
||||
@@ -83,6 +83,7 @@ func (d *HubSpotDriver) ListAccounts(ctx context.Context) ([]AccountRecord, erro
|
||||
|
||||
for _, u := range resp.Results {
|
||||
role := "User"
|
||||
|
||||
if roleMap != nil && u.RoleID != "" {
|
||||
if name, ok := roleMap[u.RoleID]; ok {
|
||||
role = name
|
||||
@@ -114,6 +115,7 @@ func (d *HubSpotDriver) ListAccounts(ctx context.Context) ([]AccountRecord, erro
|
||||
if resp.Paging == nil || resp.Paging.Next == nil || resp.Paging.Next.After == "" {
|
||||
return records, nil
|
||||
}
|
||||
|
||||
after = resp.Paging.Next.After
|
||||
}
|
||||
|
||||
@@ -128,9 +130,11 @@ func (d *HubSpotDriver) fetchUsers(ctx context.Context, after string) (*hubspotU
|
||||
|
||||
q := req.URL.Query()
|
||||
q.Set("limit", "100")
|
||||
|
||||
if after != "" {
|
||||
q.Set("after", after)
|
||||
}
|
||||
|
||||
req.URL.RawQuery = q.Encode()
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
@@ -139,6 +143,7 @@ func (d *HubSpotDriver) fetchUsers(ctx context.Context, after string) (*hubspotU
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute hubspot users request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
@@ -167,6 +172,7 @@ func (d *HubSpotDriver) fetchRoles(ctx context.Context) (map[string]string, erro
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute hubspot roles request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
|
||||
@@ -61,6 +61,7 @@ func (d *IntercomDriver) ListAccounts(ctx context.Context) ([]AccountRecord, err
|
||||
}
|
||||
|
||||
var records []AccountRecord
|
||||
|
||||
for _, a := range resp.Admins {
|
||||
record := AccountRecord{
|
||||
Email: a.Email,
|
||||
@@ -95,6 +96,7 @@ func (d *IntercomDriver) fetchAdmins(ctx context.Context) (*intercomAdminsRespon
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute intercom admins request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
@@ -119,5 +121,6 @@ func intercomRole(hasInboxSeat bool) string {
|
||||
if hasInboxSeat {
|
||||
return "Agent"
|
||||
}
|
||||
|
||||
return "Viewer"
|
||||
}
|
||||
|
||||
@@ -125,6 +125,7 @@ func (d *LinearDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error
|
||||
if !resp.Data.Users.PageInfo.HasNextPage || resp.Data.Users.PageInfo.EndCursor == "" {
|
||||
return records, nil
|
||||
}
|
||||
|
||||
nextCursor := resp.Data.Users.PageInfo.EndCursor
|
||||
after = &nextCursor
|
||||
}
|
||||
@@ -170,6 +171,7 @@ query AccessReviewLinearUsers($after: String) {
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create linear users request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
@@ -177,6 +179,7 @@ query AccessReviewLinearUsers($after: String) {
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute linear users request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
@@ -189,6 +192,7 @@ query AccessReviewLinearUsers($after: String) {
|
||||
if err := json.NewDecoder(httpResp.Body).Decode(&resp); err != nil {
|
||||
return nil, fmt.Errorf("cannot decode linear users response: %w", err)
|
||||
}
|
||||
|
||||
if len(resp.Errors) > 0 {
|
||||
return nil, fmt.Errorf("linear graphql error: %s", resp.Errors[0].Message)
|
||||
}
|
||||
|
||||
@@ -120,15 +120,18 @@ func (d *Microsoft365Driver) ListAccounts(ctx context.Context) ([]AccountRecord,
|
||||
}
|
||||
|
||||
rolesByUser := make(map[string][]string)
|
||||
|
||||
for _, role := range roles {
|
||||
members, err := d.listRoleMembers(ctx, role.ID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot list members of role %q: %w", role.DisplayName, err)
|
||||
}
|
||||
|
||||
for _, m := range members {
|
||||
if m.ODataType != "" && m.ODataType != "#microsoft.graph.user" {
|
||||
continue
|
||||
}
|
||||
|
||||
rolesByUser[m.ID] = append(rolesByUser[m.ID], role.DisplayName)
|
||||
}
|
||||
}
|
||||
@@ -147,6 +150,7 @@ func (d *Microsoft365Driver) ListAccounts(ctx context.Context) ([]AccountRecord,
|
||||
|
||||
userRoles := rolesByUser[u.ID]
|
||||
isAdmin := false
|
||||
|
||||
for _, r := range userRoles {
|
||||
if adminRoleDisplayNames[r] {
|
||||
isAdmin = true
|
||||
@@ -214,6 +218,7 @@ func pickHighestRole(roles []string) string {
|
||||
if len(roles) > 0 {
|
||||
return roles[0]
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
@@ -224,15 +229,18 @@ func (d *Microsoft365Driver) listUsers(ctx context.Context) ([]microsoft365User,
|
||||
}
|
||||
|
||||
var all []microsoft365User
|
||||
|
||||
for range microsoft365MaxPaginationOK {
|
||||
var page microsoft365UsersPage
|
||||
if err := d.fetchJSON(ctx, pageURL, &page); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
all = append(all, page.Value...)
|
||||
if page.NextLink == "" {
|
||||
return all, nil
|
||||
}
|
||||
|
||||
pageURL = page.NextLink
|
||||
}
|
||||
|
||||
@@ -258,15 +266,18 @@ func (d *Microsoft365Driver) listDirectoryRoles(ctx context.Context) ([]microsof
|
||||
url := fmt.Sprintf("%s/directoryRoles", microsoft365GraphBaseURL)
|
||||
|
||||
var all []microsoft365DirectoryRole
|
||||
|
||||
for range microsoft365MaxPaginationOK {
|
||||
var page microsoft365RolesPage
|
||||
if err := d.fetchJSON(ctx, url, &page); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
all = append(all, page.Value...)
|
||||
if page.NextLink == "" {
|
||||
return all, nil
|
||||
}
|
||||
|
||||
url = page.NextLink
|
||||
}
|
||||
|
||||
@@ -277,15 +288,18 @@ func (d *Microsoft365Driver) listRoleMembers(ctx context.Context, roleID string)
|
||||
url := fmt.Sprintf("%s/directoryRoles/%s/members", microsoft365GraphBaseURL, roleID)
|
||||
|
||||
var all []microsoft365RoleMember
|
||||
|
||||
for range microsoft365MaxPaginationOK {
|
||||
var page microsoft365MembersPage
|
||||
if err := d.fetchJSON(ctx, url, &page); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
all = append(all, page.Value...)
|
||||
if page.NextLink == "" {
|
||||
return all, nil
|
||||
}
|
||||
|
||||
url = page.NextLink
|
||||
}
|
||||
|
||||
@@ -297,12 +311,14 @@ func (d *Microsoft365Driver) fetchJSON(ctx context.Context, url string, dst any)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot create graph request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot execute graph request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
|
||||
@@ -144,6 +144,7 @@ func (d *MondayDriver) queryUsers(ctx context.Context, page int) ([]mondayUser,
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create monday users request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
@@ -151,6 +152,7 @@ func (d *MondayDriver) queryUsers(ctx context.Context, page int) ([]mondayUser,
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute monday users request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
|
||||
@@ -69,6 +69,7 @@ func ProviderDisplayName(provider coredata.ConnectorProvider) string {
|
||||
if name, ok := providerDisplayNames[provider]; ok {
|
||||
return name
|
||||
}
|
||||
|
||||
return string(provider)
|
||||
}
|
||||
|
||||
@@ -91,6 +92,7 @@ func (r *slackNameResolver) ResolveInstanceName(ctx context.Context) (string, er
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute slack auth.test request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
var resp struct {
|
||||
@@ -156,6 +158,7 @@ func (r *linearNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create linear organization request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
@@ -163,6 +166,7 @@ func (r *linearNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute linear organization request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -182,6 +186,7 @@ func (r *linearNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err := json.NewDecoder(httpResp.Body).Decode(&resp); err != nil {
|
||||
return "", fmt.Errorf("cannot decode linear organization response: %w", err)
|
||||
}
|
||||
|
||||
if len(resp.Errors) > 0 {
|
||||
return "", fmt.Errorf("linear graphql error: %s", resp.Errors[0].Message)
|
||||
}
|
||||
@@ -208,12 +213,14 @@ func (r *cloudflareNameResolver) ResolveInstanceName(ctx context.Context) (strin
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create cloudflare accounts request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute cloudflare accounts request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -255,12 +262,14 @@ func (r *brexNameResolver) ResolveInstanceName(ctx context.Context) (string, err
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create brex company request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute brex company request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -297,12 +306,14 @@ func (r *tallyNameResolver) ResolveInstanceName(ctx context.Context) (string, er
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create tally organization request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute tally organization request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -338,12 +349,14 @@ func (r *hubspotNameResolver) ResolveInstanceName(ctx context.Context) (string,
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create hubspot account-info request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute hubspot account-info request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -375,12 +388,14 @@ func (r *docusignNameResolver) ResolveInstanceName(ctx context.Context) (string,
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create docusign userinfo request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute docusign userinfo request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -429,12 +444,14 @@ func (r *openaiNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create openai organization request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute openai organization request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -473,12 +490,14 @@ func (r *sentryNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create sentry organization request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute sentry organization request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -512,12 +531,14 @@ func (r *githubNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create github organization request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/vnd.github+json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute github organization request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -565,6 +586,7 @@ func (r *intercomNameResolver) ResolveInstanceName(ctx context.Context) (string,
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create intercom me request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("Intercom-Version", "2.11")
|
||||
|
||||
@@ -572,6 +594,7 @@ func (r *intercomNameResolver) ResolveInstanceName(ctx context.Context) (string,
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute intercom me request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -622,12 +645,14 @@ func (r *gitlabNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create gitlab group request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute gitlab group request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -645,6 +670,7 @@ func (r *gitlabNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if resp.Name != "" {
|
||||
return resp.Name, nil
|
||||
}
|
||||
|
||||
return resp.FullPath, nil
|
||||
}
|
||||
|
||||
@@ -669,12 +695,14 @@ func (r *bitbucketNameResolver) ResolveInstanceName(ctx context.Context) (string
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create bitbucket workspace request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute bitbucket workspace request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -692,6 +720,7 @@ func (r *bitbucketNameResolver) ResolveInstanceName(ctx context.Context) (string
|
||||
if resp.Name != "" {
|
||||
return resp.Name, nil
|
||||
}
|
||||
|
||||
return resp.Slug, nil
|
||||
}
|
||||
|
||||
@@ -716,12 +745,14 @@ func (r *herokuNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create heroku team request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/vnd.heroku+json; version=3")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute heroku team request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -774,12 +805,14 @@ func (r *asanaNameResolver) ResolveInstanceName(ctx context.Context) (string, er
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create asana workspace request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute asana workspace request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -819,12 +852,14 @@ func (r *netlifyNameResolver) ResolveInstanceName(ctx context.Context) (string,
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create netlify account request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute netlify account request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -862,12 +897,14 @@ func (r *clickupNameResolver) ResolveInstanceName(ctx context.Context) (string,
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create clickup team request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute clickup team request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -905,16 +942,19 @@ func (r *vercelNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
}
|
||||
|
||||
teamURL := fmt.Sprintf("https://api.vercel.com/v2/teams/%s", url.PathEscape(r.teamID))
|
||||
|
||||
teamReq, err := http.NewRequestWithContext(ctx, http.MethodGet, teamURL, nil)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create vercel team request: %w", err)
|
||||
}
|
||||
|
||||
teamReq.Header.Set("Accept", "application/json")
|
||||
|
||||
teamResp, err := r.httpClient.Do(teamReq)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute vercel team request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = teamResp.Body.Close() }()
|
||||
|
||||
if teamResp.StatusCode == http.StatusOK {
|
||||
@@ -925,9 +965,11 @@ func (r *vercelNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err := json.NewDecoder(teamResp.Body).Decode(&body); err != nil {
|
||||
return "", fmt.Errorf("cannot decode vercel team response: %w", err)
|
||||
}
|
||||
|
||||
if body.Name != "" {
|
||||
return body.Name, nil
|
||||
}
|
||||
|
||||
return body.Slug, nil
|
||||
}
|
||||
|
||||
@@ -941,9 +983,11 @@ func (r *vercelNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if user.Username != "" {
|
||||
return user.Username, nil
|
||||
}
|
||||
|
||||
return user.Name, nil
|
||||
}
|
||||
|
||||
@@ -972,6 +1016,7 @@ func (r *mondayNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create monday account request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
@@ -979,6 +1024,7 @@ func (r *mondayNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute monday account request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -1023,6 +1069,7 @@ func (r *notionNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create notion users/me request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("Notion-Version", notionAPIVersion)
|
||||
|
||||
@@ -1030,6 +1077,7 @@ func (r *notionNameResolver) ResolveInstanceName(ctx context.Context) (string, e
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute notion users/me request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -1068,12 +1116,14 @@ func (r *microsoft365NameResolver) ResolveInstanceName(ctx context.Context) (str
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create microsoft 365 organization request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot execute microsoft 365 organization request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
@@ -1101,13 +1151,16 @@ func (r *microsoft365NameResolver) ResolveInstanceName(ctx context.Context) (str
|
||||
if org.DisplayName != "" {
|
||||
return org.DisplayName, nil
|
||||
}
|
||||
|
||||
for _, d := range org.VerifiedDomains {
|
||||
if d.IsDefault {
|
||||
return d.Name, nil
|
||||
}
|
||||
}
|
||||
|
||||
if len(org.VerifiedDomains) > 0 {
|
||||
return org.VerifiedDomains[0].Name, nil
|
||||
}
|
||||
|
||||
return "", nil
|
||||
}
|
||||
|
||||
@@ -37,9 +37,11 @@ func (h *hostRewriter) RoundTrip(r *http.Request) (*http.Response, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
r2 := r.Clone(r.Context())
|
||||
r2.URL.Scheme = u.Scheme
|
||||
r2.URL.Host = u.Host
|
||||
|
||||
return http.DefaultTransport.RoundTrip(r2)
|
||||
}
|
||||
|
||||
@@ -88,11 +90,13 @@ func TestNotionNameResolver(t *testing.T) {
|
||||
defer srv.Close()
|
||||
|
||||
client := &http.Client{Transport: &hostRewriter{target: srv.URL}}
|
||||
|
||||
got, err := NewNotionNameResolver(client).ResolveInstanceName(context.Background())
|
||||
if tc.wantErr {
|
||||
require.Error(t, err)
|
||||
return
|
||||
}
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.want, got)
|
||||
})
|
||||
|
||||
@@ -94,12 +94,14 @@ func (d *NetlifyDriver) queryMembers(ctx context.Context, endpoint string) ([]ne
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("cannot create netlify members request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("cannot execute netlify members request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
|
||||
@@ -96,6 +96,7 @@ func (d *NotionDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error
|
||||
if !resp.HasMore || resp.NextCursor == "" {
|
||||
return records, nil
|
||||
}
|
||||
|
||||
nextCursor := resp.NextCursor
|
||||
startCursor = &nextCursor
|
||||
}
|
||||
@@ -114,15 +115,18 @@ func (d *NotionDriver) queryUsers(ctx context.Context, startCursor *string) (*no
|
||||
|
||||
q := req.URL.Query()
|
||||
q.Set("page_size", "100")
|
||||
|
||||
if startCursor != nil {
|
||||
q.Set("start_cursor", *startCursor)
|
||||
}
|
||||
|
||||
req.URL.RawQuery = q.Encode()
|
||||
|
||||
httpResp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute notion users request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
|
||||
@@ -71,6 +71,7 @@ func NewOnePasswordDriver(httpClient *http.Client, baseURL string) *OnePasswordD
|
||||
|
||||
func (d *OnePasswordDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error) {
|
||||
var records []AccountRecord
|
||||
|
||||
startIndex := 1
|
||||
|
||||
for range maxPaginationPages {
|
||||
@@ -103,6 +104,7 @@ func (d *OnePasswordDriver) ListAccounts(ctx context.Context) ([]AccountRecord,
|
||||
if record.FullName == "" && u.Name.Formatted != "" {
|
||||
record.FullName = u.Name.Formatted
|
||||
}
|
||||
|
||||
if record.FullName == "" && (u.Name.GivenName != "" || u.Name.FamilyName != "") {
|
||||
record.FullName = u.Name.GivenName + " " + u.Name.FamilyName
|
||||
}
|
||||
@@ -128,6 +130,7 @@ func (d *OnePasswordDriver) ListAccounts(ctx context.Context) ([]AccountRecord,
|
||||
if len(resp.Resources) == 0 || resp.ItemsPerPage <= 0 || startIndex+resp.ItemsPerPage > resp.TotalResults {
|
||||
return records, nil
|
||||
}
|
||||
|
||||
startIndex += resp.ItemsPerPage
|
||||
}
|
||||
|
||||
@@ -139,6 +142,7 @@ func (d *OnePasswordDriver) queryUsers(ctx context.Context, startIndex int) (*on
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot parse 1password base url: %w", err)
|
||||
}
|
||||
|
||||
u = u.JoinPath("scim", "v2", "Users")
|
||||
q := u.Query()
|
||||
q.Set("startIndex", strconv.Itoa(startIndex))
|
||||
@@ -149,12 +153,14 @@ func (d *OnePasswordDriver) queryUsers(ctx context.Context, startIndex int) (*on
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create 1password users request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/scim+json")
|
||||
|
||||
httpResp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute 1password users request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
|
||||
@@ -107,6 +107,7 @@ func (d *OnePasswordUsersAPIDriver) ListAccounts(ctx context.Context) ([]Account
|
||||
if resp.NextPageToken == "" {
|
||||
return records, nil
|
||||
}
|
||||
|
||||
pageToken = resp.NextPageToken
|
||||
}
|
||||
|
||||
@@ -118,6 +119,7 @@ func (d *OnePasswordUsersAPIDriver) queryUsers(ctx context.Context, pageToken st
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot parse 1password users api base url: %w", err)
|
||||
}
|
||||
|
||||
u = u.JoinPath("v1beta1", "accounts", d.accountID, "users")
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
|
||||
@@ -127,9 +129,11 @@ func (d *OnePasswordUsersAPIDriver) queryUsers(ctx context.Context, pageToken st
|
||||
|
||||
q := req.URL.Query()
|
||||
q.Set("max_page_size", "100")
|
||||
|
||||
if pageToken != "" {
|
||||
q.Set("page_token", pageToken)
|
||||
}
|
||||
|
||||
req.URL.RawQuery = q.Encode()
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
@@ -138,6 +142,7 @@ func (d *OnePasswordUsersAPIDriver) queryUsers(ctx context.Context, pageToken st
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute 1password users api request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
|
||||
@@ -89,6 +89,7 @@ func (d *OpenAIDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error
|
||||
if !resp.HasMore || resp.LastID == "" {
|
||||
return records, nil
|
||||
}
|
||||
|
||||
after = resp.LastID
|
||||
}
|
||||
|
||||
@@ -103,9 +104,11 @@ func (d *OpenAIDriver) fetchUsers(ctx context.Context, after string) (*openaiUse
|
||||
|
||||
q := req.URL.Query()
|
||||
q.Set("limit", "100")
|
||||
|
||||
if after != "" {
|
||||
q.Set("after", after)
|
||||
}
|
||||
|
||||
req.URL.RawQuery = q.Encode()
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
@@ -114,6 +117,7 @@ func (d *OpenAIDriver) fetchUsers(ctx context.Context, after string) (*openaiUse
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute openai users request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
|
||||
@@ -38,12 +38,14 @@ func ListGitHubOrganizations(ctx context.Context, httpClient *http.Client) ([]Or
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create github organizations request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot fetch github organizations: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
@@ -64,8 +66,10 @@ func ListGitHubOrganizations(ctx context.Context, httpClient *http.Client) ([]Or
|
||||
if displayName == "" {
|
||||
displayName = org.Login
|
||||
}
|
||||
|
||||
result[i] = Organization{Slug: org.Login, DisplayName: displayName}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -81,12 +85,14 @@ func ListSentryOrganizations(ctx context.Context, httpClient *http.Client) ([]Or
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create sentry organizations request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot fetch sentry organizations: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
@@ -107,8 +113,10 @@ func ListSentryOrganizations(ctx context.Context, httpClient *http.Client) ([]Or
|
||||
if displayName == "" {
|
||||
displayName = org.Slug
|
||||
}
|
||||
|
||||
result[i] = Organization{Slug: org.Slug, DisplayName: displayName}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -125,12 +133,14 @@ func ListGitLabOrganizations(ctx context.Context, httpClient *http.Client) ([]Or
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create gitlab organizations request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot fetch gitlab organizations: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
@@ -152,11 +162,13 @@ func ListGitLabOrganizations(ctx context.Context, httpClient *http.Client) ([]Or
|
||||
if displayName == "" {
|
||||
displayName = g.FullPath
|
||||
}
|
||||
|
||||
result[i] = Organization{
|
||||
Slug: strconv.FormatInt(g.ID, 10),
|
||||
DisplayName: displayName,
|
||||
}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -174,6 +186,7 @@ func ListBitbucketOrganizations(ctx context.Context, httpClient *http.Client) ([
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create bitbucket organizations request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := httpClient.Do(req)
|
||||
@@ -203,6 +216,7 @@ func ListBitbucketOrganizations(ctx context.Context, httpClient *http.Client) ([
|
||||
_ = resp.Body.Close()
|
||||
return nil, fmt.Errorf("cannot decode bitbucket organizations response: %w", err)
|
||||
}
|
||||
|
||||
_ = resp.Body.Close()
|
||||
|
||||
for _, v := range body.Values {
|
||||
@@ -211,18 +225,22 @@ func ListBitbucketOrganizations(ctx context.Context, httpClient *http.Client) ([
|
||||
slug = v.Workspace.Slug
|
||||
name = v.Workspace.Name
|
||||
}
|
||||
|
||||
displayName := name
|
||||
if displayName == "" {
|
||||
displayName = slug
|
||||
}
|
||||
|
||||
result = append(result, Organization{Slug: slug, DisplayName: displayName})
|
||||
}
|
||||
|
||||
if body.Next == "" {
|
||||
return result, nil
|
||||
}
|
||||
|
||||
pageURL = body.Next
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("cannot list all bitbucket organizations: %w", ErrPaginationLimitReached)
|
||||
}
|
||||
|
||||
@@ -233,12 +251,14 @@ func ListHerokuOrganizations(ctx context.Context, httpClient *http.Client) ([]Or
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create heroku organizations request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/vnd.heroku+json; version=3")
|
||||
|
||||
resp, err := httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot fetch heroku organizations: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
@@ -259,8 +279,10 @@ func ListHerokuOrganizations(ctx context.Context, httpClient *http.Client) ([]Or
|
||||
if displayName == "" {
|
||||
displayName = t.ID
|
||||
}
|
||||
|
||||
result[i] = Organization{Slug: t.ID, DisplayName: displayName}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -276,12 +298,14 @@ func ListAsanaOrganizations(ctx context.Context, httpClient *http.Client) ([]Org
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create asana organizations request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot fetch asana organizations: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
@@ -304,8 +328,10 @@ func ListAsanaOrganizations(ctx context.Context, httpClient *http.Client) ([]Org
|
||||
if displayName == "" {
|
||||
displayName = w.GID
|
||||
}
|
||||
|
||||
result[i] = Organization{Slug: w.GID, DisplayName: displayName}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -321,12 +347,14 @@ func ListNetlifyOrganizations(ctx context.Context, httpClient *http.Client) ([]O
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create netlify organizations request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot fetch netlify organizations: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
@@ -348,8 +376,10 @@ func ListNetlifyOrganizations(ctx context.Context, httpClient *http.Client) ([]O
|
||||
if displayName == "" {
|
||||
displayName = a.Slug
|
||||
}
|
||||
|
||||
result[i] = Organization{Slug: a.Slug, DisplayName: displayName}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -365,12 +395,14 @@ func ListClickUpOrganizations(ctx context.Context, httpClient *http.Client) ([]O
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create clickup organizations request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot fetch clickup organizations: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
@@ -393,7 +425,9 @@ func ListClickUpOrganizations(ctx context.Context, httpClient *http.Client) ([]O
|
||||
if displayName == "" {
|
||||
displayName = t.ID
|
||||
}
|
||||
|
||||
result[i] = Organization{Slug: t.ID, DisplayName: displayName}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -65,6 +65,7 @@ func (d *PagerDutyDriver) ListAccounts(ctx context.Context) ([]AccountRecord, er
|
||||
var records []AccountRecord
|
||||
|
||||
const limit = 100
|
||||
|
||||
offset := 0
|
||||
|
||||
for range maxPaginationPages {
|
||||
@@ -113,6 +114,7 @@ func (d *PagerDutyDriver) ListAccounts(ctx context.Context) ([]AccountRecord, er
|
||||
if pageSize <= 0 {
|
||||
pageSize = limit
|
||||
}
|
||||
|
||||
offset += pageSize
|
||||
}
|
||||
|
||||
@@ -130,12 +132,14 @@ func (d *PagerDutyDriver) queryUsers(ctx context.Context, offset, limit int) (*p
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create pagerduty users request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/vnd.pagerduty+json;version=2")
|
||||
|
||||
httpResp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute pagerduty users request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
|
||||
@@ -54,6 +54,7 @@ func (d *ResendDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error
|
||||
}
|
||||
|
||||
var records []AccountRecord
|
||||
|
||||
for _, k := range resp.Data {
|
||||
record := AccountRecord{
|
||||
FullName: k.Name,
|
||||
@@ -96,6 +97,7 @@ func (d *ResendDriver) fetchAPIKeys(ctx context.Context) (*resendAPIKeysResponse
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute resend api-keys request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
|
||||
@@ -67,12 +67,14 @@ func (d *SentryDriver) resolveOrgSlug(ctx context.Context) (string, error) {
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot create sentry organizations request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot fetch sentry organizations: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
@@ -100,6 +102,7 @@ func (d *SentryDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot resolve sentry organization slug: %w", err)
|
||||
}
|
||||
|
||||
orgSlug = slug
|
||||
}
|
||||
|
||||
@@ -130,6 +133,7 @@ func (d *SentryDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error
|
||||
isAdmin := m.OrgRole == "admin" || m.OrgRole == "owner"
|
||||
|
||||
mfaStatus := coredata.MFAStatusUnknown
|
||||
|
||||
if m.User != nil {
|
||||
if m.User.Has2FA {
|
||||
mfaStatus = coredata.MFAStatusEnabled
|
||||
@@ -188,6 +192,7 @@ func (d *SentryDriver) queryMembers(ctx context.Context, url string) ([]sentryMe
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("cannot execute sentry members request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
@@ -221,8 +226,10 @@ func sentryAuthMethod(flags map[string]bool, user *sentryUser) coredata.AccessEn
|
||||
if flags["sso:linked"] {
|
||||
return coredata.AccessEntryAuthMethodSSO
|
||||
}
|
||||
|
||||
if user != nil && user.HasPasswordAuth {
|
||||
return coredata.AccessEntryAuthMethodPassword
|
||||
}
|
||||
|
||||
return coredata.AccessEntryAuthMethodUnknown
|
||||
}
|
||||
|
||||
@@ -120,6 +120,7 @@ func (d *SlackDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error)
|
||||
if resp.ResponseMetadata.NextCursor == "" {
|
||||
return records, nil
|
||||
}
|
||||
|
||||
cursor = resp.ResponseMetadata.NextCursor
|
||||
}
|
||||
|
||||
@@ -134,15 +135,18 @@ func (d *SlackDriver) queryUsers(ctx context.Context, cursor string) (*slackUser
|
||||
|
||||
q := req.URL.Query()
|
||||
q.Set("limit", "200")
|
||||
|
||||
if cursor != "" {
|
||||
q.Set("cursor", cursor)
|
||||
}
|
||||
|
||||
req.URL.RawQuery = q.Encode()
|
||||
|
||||
httpResp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute slack users.list request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
@@ -180,5 +184,6 @@ func slackMFAStatus(has2FA bool) coredata.MFAStatus {
|
||||
if has2FA {
|
||||
return coredata.MFAStatusEnabled
|
||||
}
|
||||
|
||||
return coredata.MFAStatusDisabled
|
||||
}
|
||||
|
||||
@@ -36,12 +36,14 @@ func TestSlackDriver(t *testing.T) {
|
||||
|
||||
// Find the first human user (bots may not have email).
|
||||
var r AccountRecord
|
||||
|
||||
for _, rec := range records {
|
||||
if rec.Email != "" {
|
||||
r = rec
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
require.NotEmpty(t, r.Email, "expected at least one record with an email")
|
||||
assert.NotEmpty(t, r.ExternalID)
|
||||
assert.NotEmpty(t, r.Role)
|
||||
|
||||
@@ -53,6 +53,7 @@ func (d *SupabaseDriver) ListAccounts(ctx context.Context) ([]AccountRecord, err
|
||||
}
|
||||
|
||||
var records []AccountRecord
|
||||
|
||||
for _, m := range members {
|
||||
mfaStatus := coredata.MFAStatusDisabled
|
||||
if m.MFAEnabled {
|
||||
@@ -96,6 +97,7 @@ func (d *SupabaseDriver) queryMembers(ctx context.Context) ([]supabaseMember, er
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute supabase members request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
|
||||
@@ -89,6 +89,7 @@ func (d *TallyDriver) listUsers(ctx context.Context) ([]AccountRecord, error) {
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute tally users request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
@@ -106,6 +107,7 @@ func (d *TallyDriver) listUsers(ctx context.Context) ([]AccountRecord, error) {
|
||||
}
|
||||
|
||||
var records []AccountRecord
|
||||
|
||||
for _, u := range users {
|
||||
mfaStatus := coredata.MFAStatusDisabled
|
||||
if u.HasTwoFactorEnabled {
|
||||
@@ -149,6 +151,7 @@ func (d *TallyDriver) listInvites(ctx context.Context) ([]AccountRecord, error)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute tally invites request: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = httpResp.Body.Close()
|
||||
}()
|
||||
@@ -166,6 +169,7 @@ func (d *TallyDriver) listInvites(ctx context.Context) ([]AccountRecord, error)
|
||||
}
|
||||
|
||||
var records []AccountRecord
|
||||
|
||||
for _, inv := range invites {
|
||||
record := AccountRecord{
|
||||
Email: inv.Email,
|
||||
|
||||
@@ -57,6 +57,7 @@ func newRecorder(t *testing.T, cassettePath string, envVar string) *recorder.Rec
|
||||
if mode == recorder.ModeReplayOnly {
|
||||
t.Skipf("cassette not found (record with %s env var): %v", envVar, err)
|
||||
}
|
||||
|
||||
t.Fatalf("cannot create vcr recorder: %v", err)
|
||||
}
|
||||
|
||||
@@ -81,6 +82,7 @@ func (rt *authRoundTripper) RoundTrip(req *http.Request) (*http.Response, error)
|
||||
if rt.authValue != "" {
|
||||
req.Header.Set("Authorization", rt.authValue)
|
||||
}
|
||||
|
||||
return rt.transport.RoundTrip(req)
|
||||
}
|
||||
|
||||
@@ -89,6 +91,7 @@ func bearerAuth(token string) string {
|
||||
if token == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
return "Bearer " + token
|
||||
}
|
||||
|
||||
@@ -104,5 +107,6 @@ func newVCRClient(rec *recorder.Recorder, authValue string) *http.Client {
|
||||
transport: transport,
|
||||
}
|
||||
}
|
||||
|
||||
return &http.Client{Transport: transport}
|
||||
}
|
||||
|
||||
@@ -111,6 +111,7 @@ func (d *VercelDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error
|
||||
if page.Pagination.Next == nil {
|
||||
return records, nil
|
||||
}
|
||||
|
||||
cursor = strconv.FormatInt(*page.Pagination.Next, 10)
|
||||
}
|
||||
|
||||
@@ -120,9 +121,11 @@ func (d *VercelDriver) ListAccounts(ctx context.Context) ([]AccountRecord, error
|
||||
func (d *VercelDriver) queryMembers(ctx context.Context, cursor string) (*vercelMembersPage, error) {
|
||||
q := url.Values{}
|
||||
q.Set("limit", "100")
|
||||
|
||||
if cursor != "" {
|
||||
q.Set("until", cursor)
|
||||
}
|
||||
|
||||
u := url.URL{
|
||||
Scheme: "https",
|
||||
Host: "api.vercel.com",
|
||||
@@ -134,12 +137,14 @@ func (d *VercelDriver) queryMembers(ctx context.Context, cursor string) (*vercel
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create vercel members request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
httpResp, err := d.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute vercel members request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = httpResp.Body.Close() }()
|
||||
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
|
||||
@@ -80,11 +80,13 @@ func (e *ReviewEngine) FetchSource(
|
||||
if err := source.LoadByID(ctx, tx, e.scope, sourceID); err != nil {
|
||||
return fmt.Errorf("cannot load access source %s: %w", sourceID, err)
|
||||
}
|
||||
|
||||
if source.OrganizationID != campaign.OrganizationID {
|
||||
return fmt.Errorf("cannot process access source: %s does not belong to campaign organization", sourceID)
|
||||
}
|
||||
|
||||
var err error
|
||||
|
||||
driver, err = e.resolveDriver(ctx, tx, source)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot resolve driver for source %s: %w", source.Name, err)
|
||||
@@ -97,6 +99,7 @@ func (e *ReviewEngine) FetchSource(
|
||||
}
|
||||
} else {
|
||||
entries := &coredata.AccessEntries{}
|
||||
|
||||
baseline, err = entries.LoadBaselineBySourceID(ctx, tx, e.scope, lastCompletedCampaign.ID, sourceID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot load baseline entries by source: %w", err)
|
||||
@@ -117,10 +120,13 @@ func (e *ReviewEngine) FetchSource(
|
||||
|
||||
sourceCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
accounts, err := driver.ListAccounts(sourceCtx)
|
||||
|
||||
cancel()
|
||||
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("cannot list accounts from source %s: %w", source.Name, err)
|
||||
}
|
||||
|
||||
fetchedCount = len(accounts)
|
||||
|
||||
err = e.pg.WithTx(
|
||||
@@ -132,6 +138,7 @@ func (e *ReviewEngine) FetchSource(
|
||||
for _, account := range accounts {
|
||||
accountKey := normalizeAccountKey(account.Email, account.ExternalID)
|
||||
seenAccountKeys[accountKey] = struct{}{}
|
||||
|
||||
incrementalTag := coredata.AccessEntryIncrementalTagNew
|
||||
if _, ok := previousByAccountKey[accountKey]; ok {
|
||||
incrementalTag = coredata.AccessEntryIncrementalTagUnchanged
|
||||
@@ -210,6 +217,7 @@ func (e *ReviewEngine) FetchSource(
|
||||
|
||||
func normalizeAccountKey(email, externalID string) string {
|
||||
emailKey := strings.ToLower(strings.TrimSpace(email))
|
||||
|
||||
externalID = strings.TrimSpace(externalID)
|
||||
if externalID != "" {
|
||||
return emailKey + "|" + externalID
|
||||
@@ -231,6 +239,7 @@ func (e *ReviewEngine) oauthClient(
|
||||
return conn.RefreshableClient(ctx, *refreshCfg)
|
||||
}
|
||||
}
|
||||
|
||||
return conn.Client(ctx)
|
||||
}
|
||||
|
||||
@@ -245,6 +254,7 @@ func (e *ReviewEngine) connectorHTTPClient(
|
||||
if oauth2Conn, ok := dbConnector.Connection.(*connector.OAuth2Connection); ok {
|
||||
return e.oauthClient(ctx, oauth2Conn, dbConnector.Provider)
|
||||
}
|
||||
|
||||
return dbConnector.Connection.Client(ctx)
|
||||
}
|
||||
|
||||
@@ -312,15 +322,19 @@ func (e *ReviewEngine) resolveDriver(
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read 1password users api settings: %w", err)
|
||||
}
|
||||
|
||||
return drivers.NewOnePasswordUsersAPIDriver(httpClient, settings.AccountID, settings.Region), nil
|
||||
}
|
||||
|
||||
onePasswordSettings, err := coredata.ConnectorSettings[coredata.OnePasswordConnectorSettings](dbConnector)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read 1password connector settings: %w", err)
|
||||
}
|
||||
|
||||
if onePasswordSettings.SCIMBridgeURL == "" {
|
||||
return nil, fmt.Errorf("1password connector requires scim_bridge_url in settings")
|
||||
}
|
||||
|
||||
return drivers.NewOnePasswordDriver(httpClient, onePasswordSettings.SCIMBridgeURL), nil
|
||||
case coredata.ConnectorProviderHubSpot:
|
||||
return drivers.NewHubSpotDriver(httpClient), nil
|
||||
@@ -335,9 +349,11 @@ func (e *ReviewEngine) resolveDriver(
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read tally connector settings: %w", err)
|
||||
}
|
||||
|
||||
if tallySettings.OrganizationID == "" {
|
||||
return nil, fmt.Errorf("tally connector requires organization_id in settings")
|
||||
}
|
||||
|
||||
return drivers.NewTallyDriver(httpClient, tallySettings.OrganizationID), nil
|
||||
case coredata.ConnectorProviderCloudflare:
|
||||
return drivers.NewCloudflareDriver(httpClient), nil
|
||||
@@ -348,6 +364,7 @@ func (e *ReviewEngine) resolveDriver(
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read sentry connector settings: %w", err)
|
||||
}
|
||||
|
||||
// OrganizationSlug may be empty for OAuth connections; the driver auto-discovers it.
|
||||
return drivers.NewSentryDriver(httpClient, sentrySettings.OrganizationSlug), nil
|
||||
case coredata.ConnectorProviderSupabase:
|
||||
@@ -355,18 +372,22 @@ func (e *ReviewEngine) resolveDriver(
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read supabase connector settings: %w", err)
|
||||
}
|
||||
|
||||
if supabaseSettings.OrganizationSlug == "" {
|
||||
return nil, fmt.Errorf("supabase connector requires organization_slug in settings")
|
||||
}
|
||||
|
||||
return drivers.NewSupabaseDriver(httpClient, supabaseSettings.OrganizationSlug), nil
|
||||
case coredata.ConnectorProviderGitHub:
|
||||
githubSettings, err := coredata.ConnectorSettings[coredata.GitHubConnectorSettings](dbConnector)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read github connector settings: %w", err)
|
||||
}
|
||||
|
||||
if githubSettings.Organization == "" {
|
||||
return nil, fmt.Errorf("github connector requires organization in settings")
|
||||
}
|
||||
|
||||
return drivers.NewGitHubDriver(httpClient, githubSettings.Organization, e.logger.Named("github")), nil
|
||||
case coredata.ConnectorProviderIntercom:
|
||||
return drivers.NewIntercomDriver(httpClient), nil
|
||||
@@ -379,27 +400,33 @@ func (e *ReviewEngine) resolveDriver(
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read gitlab connector settings: %w", err)
|
||||
}
|
||||
|
||||
if gitlabSettings.GroupID == "" {
|
||||
return nil, fmt.Errorf("gitlab connector requires group_id in settings")
|
||||
}
|
||||
|
||||
return drivers.NewGitLabDriver(httpClient, gitlabSettings.GroupID), nil
|
||||
case coredata.ConnectorProviderBitbucket:
|
||||
bitbucketSettings, err := coredata.ConnectorSettings[coredata.BitbucketConnectorSettings](dbConnector)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read bitbucket connector settings: %w", err)
|
||||
}
|
||||
|
||||
if bitbucketSettings.Workspace == "" {
|
||||
return nil, fmt.Errorf("bitbucket connector requires workspace in settings")
|
||||
}
|
||||
|
||||
return drivers.NewBitbucketDriver(httpClient, bitbucketSettings.Workspace), nil
|
||||
case coredata.ConnectorProviderHeroku:
|
||||
herokuSettings, err := coredata.ConnectorSettings[coredata.HerokuConnectorSettings](dbConnector)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read heroku connector settings: %w", err)
|
||||
}
|
||||
|
||||
if herokuSettings.TeamID == "" {
|
||||
return nil, fmt.Errorf("heroku connector requires team_id in settings")
|
||||
}
|
||||
|
||||
return drivers.NewHerokuDriver(httpClient, herokuSettings.TeamID), nil
|
||||
case coredata.ConnectorProviderPagerDuty:
|
||||
// PagerDuty's REST API uses the regional api.pagerduty.com host;
|
||||
@@ -413,36 +440,44 @@ func (e *ReviewEngine) resolveDriver(
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read asana connector settings: %w", err)
|
||||
}
|
||||
|
||||
if asanaSettings.WorkspaceGID == "" {
|
||||
return nil, fmt.Errorf("asana connector requires workspace_gid in settings")
|
||||
}
|
||||
|
||||
return drivers.NewAsanaDriver(httpClient, asanaSettings.WorkspaceGID), nil
|
||||
case coredata.ConnectorProviderNetlify:
|
||||
netlifySettings, err := coredata.ConnectorSettings[coredata.NetlifyConnectorSettings](dbConnector)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read netlify connector settings: %w", err)
|
||||
}
|
||||
|
||||
if netlifySettings.AccountSlug == "" {
|
||||
return nil, fmt.Errorf("netlify connector requires account_slug in settings")
|
||||
}
|
||||
|
||||
return drivers.NewNetlifyDriver(httpClient, netlifySettings.AccountSlug), nil
|
||||
case coredata.ConnectorProviderClickUp:
|
||||
clickupSettings, err := coredata.ConnectorSettings[coredata.ClickUpConnectorSettings](dbConnector)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read clickup connector settings: %w", err)
|
||||
}
|
||||
|
||||
if clickupSettings.TeamID == "" {
|
||||
return nil, fmt.Errorf("clickup connector requires team_id in settings")
|
||||
}
|
||||
|
||||
return drivers.NewClickUpDriver(httpClient, clickupSettings.TeamID), nil
|
||||
case coredata.ConnectorProviderVercel:
|
||||
vercelSettings, err := coredata.ConnectorSettings[coredata.VercelConnectorSettings](dbConnector)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read vercel connector settings: %w", err)
|
||||
}
|
||||
|
||||
if vercelSettings.TeamID == "" {
|
||||
return nil, fmt.Errorf("vercel connector requires team_id in settings")
|
||||
}
|
||||
|
||||
return drivers.NewVercelDriver(httpClient, vercelSettings.TeamID), nil
|
||||
case coredata.ConnectorProviderMonday:
|
||||
return drivers.NewMondayDriver(httpClient), nil
|
||||
|
||||
@@ -78,6 +78,7 @@ func NewService(
|
||||
} else {
|
||||
fetchWorkerOpts = append(fetchWorkerOpts, worker.WithInterval(30*time.Second))
|
||||
}
|
||||
|
||||
fetchWorkerOpts = append(fetchWorkerOpts, worker.WithMaxConcurrency(20))
|
||||
|
||||
s.fetchWorker = NewSourceFetchWorker(
|
||||
@@ -137,11 +138,14 @@ func (s *Service) ResolveEntryOrganizationID(ctx context.Context, entryID gid.GI
|
||||
ctx,
|
||||
func(ctx context.Context, conn pg.Querier) error {
|
||||
var err error
|
||||
|
||||
entry := &coredata.AccessEntry{}
|
||||
|
||||
organizationID, err = entry.LoadOrganizationID(ctx, conn, entryID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot load organization id: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
)
|
||||
|
||||
@@ -79,6 +79,7 @@ func (h *sourceNameHandler) Claim(ctx context.Context) (coredata.AccessSource, e
|
||||
if errors.Is(err, coredata.ErrNoAccessSourceNameSyncAvailable) {
|
||||
return coredata.AccessSource{}, worker.ErrNoTask
|
||||
}
|
||||
|
||||
return coredata.AccessSource{}, err
|
||||
}
|
||||
|
||||
@@ -128,6 +129,7 @@ func (h *sourceNameHandler) Process(ctx context.Context, source coredata.AccessS
|
||||
}
|
||||
|
||||
resolver = h.buildResolver(&dbConnector, httpClient)
|
||||
|
||||
return nil
|
||||
},
|
||||
)
|
||||
@@ -136,6 +138,7 @@ func (h *sourceNameHandler) Process(ctx context.Context, source coredata.AccessS
|
||||
log.String("source_id", source.ID.String()),
|
||||
log.Error(err),
|
||||
)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -144,6 +147,7 @@ func (h *sourceNameHandler) Process(ctx context.Context, source coredata.AccessS
|
||||
log.String("source_id", source.ID.String()),
|
||||
log.String("provider", dbConnector.Provider.String()),
|
||||
)
|
||||
|
||||
return h.markNameSynced(ctx, &source)
|
||||
}
|
||||
|
||||
@@ -157,6 +161,7 @@ func (h *sourceNameHandler) Process(ctx context.Context, source coredata.AccessS
|
||||
log.String("provider", dbConnector.Provider.String()),
|
||||
log.Error(err),
|
||||
)
|
||||
|
||||
return fmt.Errorf("cannot resolve instance name for source %s: %w", source.ID, err)
|
||||
}
|
||||
|
||||
@@ -165,6 +170,7 @@ func (h *sourceNameHandler) Process(ctx context.Context, source coredata.AccessS
|
||||
log.String("source_id", source.ID.String()),
|
||||
log.String("provider", dbConnector.Provider.String()),
|
||||
)
|
||||
|
||||
return h.markNameSynced(ctx, &source)
|
||||
}
|
||||
|
||||
@@ -178,6 +184,7 @@ func (h *sourceNameHandler) Process(ctx context.Context, source coredata.AccessS
|
||||
)
|
||||
|
||||
source.Name = newName
|
||||
|
||||
return h.markNameSynced(ctx, &source)
|
||||
}
|
||||
|
||||
@@ -247,6 +254,7 @@ func (h *sourceNameHandler) buildResolver(
|
||||
h.logger.Error("cannot read tally connector settings", log.Error(err))
|
||||
return nil
|
||||
}
|
||||
|
||||
return drivers.NewTallyNameResolver(httpClient, tallySettings.OrganizationID)
|
||||
case coredata.ConnectorProviderHubSpot:
|
||||
return drivers.NewHubSpotNameResolver(httpClient)
|
||||
@@ -260,6 +268,7 @@ func (h *sourceNameHandler) buildResolver(
|
||||
h.logger.Error("cannot read sentry connector settings", log.Error(err))
|
||||
return nil
|
||||
}
|
||||
|
||||
return drivers.NewSentryNameResolver(httpClient, sentrySettings.OrganizationSlug)
|
||||
case coredata.ConnectorProviderGitHub:
|
||||
githubSettings, err := coredata.ConnectorSettings[coredata.GitHubConnectorSettings](dbConnector)
|
||||
@@ -267,6 +276,7 @@ func (h *sourceNameHandler) buildResolver(
|
||||
h.logger.Error("cannot read github connector settings", log.Error(err))
|
||||
return nil
|
||||
}
|
||||
|
||||
return drivers.NewGitHubNameResolver(httpClient, githubSettings.Organization)
|
||||
case coredata.ConnectorProviderSupabase:
|
||||
supabaseSettings, err := coredata.ConnectorSettings[coredata.SupabaseConnectorSettings](dbConnector)
|
||||
@@ -274,6 +284,7 @@ func (h *sourceNameHandler) buildResolver(
|
||||
h.logger.Error("cannot read supabase connector settings", log.Error(err))
|
||||
return nil
|
||||
}
|
||||
|
||||
return drivers.NewSupabaseNameResolver(supabaseSettings.OrganizationSlug)
|
||||
case coredata.ConnectorProviderIntercom:
|
||||
return drivers.NewIntercomNameResolver(httpClient)
|
||||
@@ -289,6 +300,7 @@ func (h *sourceNameHandler) buildResolver(
|
||||
h.logger.Error("cannot read gitlab connector settings", log.Error(err))
|
||||
return nil
|
||||
}
|
||||
|
||||
return drivers.NewGitLabNameResolver(httpClient, gitlabSettings.GroupID)
|
||||
case coredata.ConnectorProviderBitbucket:
|
||||
bitbucketSettings, err := coredata.ConnectorSettings[coredata.BitbucketConnectorSettings](dbConnector)
|
||||
@@ -296,6 +308,7 @@ func (h *sourceNameHandler) buildResolver(
|
||||
h.logger.Error("cannot read bitbucket connector settings", log.Error(err))
|
||||
return nil
|
||||
}
|
||||
|
||||
return drivers.NewBitbucketNameResolver(httpClient, bitbucketSettings.Workspace)
|
||||
case coredata.ConnectorProviderHeroku:
|
||||
herokuSettings, err := coredata.ConnectorSettings[coredata.HerokuConnectorSettings](dbConnector)
|
||||
@@ -303,6 +316,7 @@ func (h *sourceNameHandler) buildResolver(
|
||||
h.logger.Error("cannot read heroku connector settings", log.Error(err))
|
||||
return nil
|
||||
}
|
||||
|
||||
return drivers.NewHerokuNameResolver(httpClient, herokuSettings.TeamID)
|
||||
case coredata.ConnectorProviderPagerDuty:
|
||||
pdSettings, err := coredata.ConnectorSettings[coredata.PagerDutyConnectorSettings](dbConnector)
|
||||
@@ -310,6 +324,7 @@ func (h *sourceNameHandler) buildResolver(
|
||||
h.logger.Error("cannot read pagerduty connector settings", log.Error(err))
|
||||
return nil
|
||||
}
|
||||
|
||||
return drivers.NewPagerDutyNameResolver(pdSettings.Subdomain)
|
||||
case coredata.ConnectorProviderAsana:
|
||||
asanaSettings, err := coredata.ConnectorSettings[coredata.AsanaConnectorSettings](dbConnector)
|
||||
@@ -317,6 +332,7 @@ func (h *sourceNameHandler) buildResolver(
|
||||
h.logger.Error("cannot read asana connector settings", log.Error(err))
|
||||
return nil
|
||||
}
|
||||
|
||||
return drivers.NewAsanaNameResolver(httpClient, asanaSettings.WorkspaceGID)
|
||||
case coredata.ConnectorProviderNetlify:
|
||||
netlifySettings, err := coredata.ConnectorSettings[coredata.NetlifyConnectorSettings](dbConnector)
|
||||
@@ -324,6 +340,7 @@ func (h *sourceNameHandler) buildResolver(
|
||||
h.logger.Error("cannot read netlify connector settings", log.Error(err))
|
||||
return nil
|
||||
}
|
||||
|
||||
return drivers.NewNetlifyNameResolver(httpClient, netlifySettings.AccountSlug)
|
||||
case coredata.ConnectorProviderClickUp:
|
||||
clickupSettings, err := coredata.ConnectorSettings[coredata.ClickUpConnectorSettings](dbConnector)
|
||||
@@ -331,6 +348,7 @@ func (h *sourceNameHandler) buildResolver(
|
||||
h.logger.Error("cannot read clickup connector settings", log.Error(err))
|
||||
return nil
|
||||
}
|
||||
|
||||
return drivers.NewClickUpNameResolver(httpClient, clickupSettings.TeamID)
|
||||
case coredata.ConnectorProviderVercel:
|
||||
vercelSettings, err := coredata.ConnectorSettings[coredata.VercelConnectorSettings](dbConnector)
|
||||
@@ -338,6 +356,7 @@ func (h *sourceNameHandler) buildResolver(
|
||||
h.logger.Error("cannot read vercel connector settings", log.Error(err))
|
||||
return nil
|
||||
}
|
||||
|
||||
return drivers.NewVercelNameResolver(httpClient, vercelSettings.TeamID)
|
||||
case coredata.ConnectorProviderMonday:
|
||||
return drivers.NewMondayNameResolver(httpClient)
|
||||
|
||||
@@ -77,12 +77,14 @@ func (h *sourceFetchHandler) Claim(ctx context.Context) (coredata.AccessReviewCa
|
||||
if err := sourceFetch.Update(ctx, tx, scope); err != nil {
|
||||
return fmt.Errorf("cannot update source fetch status: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
); err != nil {
|
||||
if errors.Is(err, coredata.ErrNoAccessReviewCampaignSourceFetchAvailable) {
|
||||
return coredata.AccessReviewCampaignSourceFetch{}, worker.ErrNoTask
|
||||
}
|
||||
|
||||
return coredata.AccessReviewCampaignSourceFetch{}, fmt.Errorf("cannot claim source fetch: %w", err)
|
||||
}
|
||||
|
||||
@@ -101,6 +103,7 @@ func (h *sourceFetchHandler) RecoverStale(ctx context.Context) error {
|
||||
ctx,
|
||||
func(ctx context.Context, tx pg.Tx) error {
|
||||
var fetches coredata.AccessReviewCampaignSourceFetches
|
||||
|
||||
count, err := fetches.RecoverStale(ctx, tx, staleThreshold, now)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot recover stale source fetches: %w", err)
|
||||
@@ -135,6 +138,7 @@ func (h *sourceFetchHandler) handle(
|
||||
if commitErr != nil {
|
||||
return fmt.Errorf("cannot load campaign: %w, and cannot commit failed source fetch: %w", err, commitErr)
|
||||
}
|
||||
|
||||
return fmt.Errorf("cannot load campaign: %w", err)
|
||||
}
|
||||
|
||||
@@ -148,6 +152,7 @@ func (h *sourceFetchHandler) handle(
|
||||
if finalizeErr := h.finalizeCampaignFetchLifecycle(ctx, sourceFetch.TenantID, sourceFetch.AccessReviewCampaignID); finalizeErr != nil {
|
||||
return fmt.Errorf("cannot finalize campaign after failed source fetch: %w", finalizeErr)
|
||||
}
|
||||
|
||||
return fmt.Errorf("cannot fetch source: %w", err)
|
||||
}
|
||||
|
||||
@@ -250,6 +255,7 @@ func (h *sourceFetchHandler) finalizeCampaignFetchLifecycle(
|
||||
|
||||
campaign.Status = coredata.AccessReviewCampaignStatusPendingActions
|
||||
campaign.UpdatedAt = time.Now()
|
||||
|
||||
return campaign.Update(ctx, tx, scope)
|
||||
},
|
||||
)
|
||||
|
||||
@@ -124,6 +124,7 @@ func (a *Agent) Clone(opts ...Option) *Agent {
|
||||
copy(newApproval.ToolNames, a.approval.ToolNames)
|
||||
newApproval.toolNameSet = buildToolNameSet(newApproval.ToolNames)
|
||||
}
|
||||
|
||||
cp.approval = &newApproval
|
||||
}
|
||||
|
||||
@@ -205,6 +206,7 @@ func WithMaxTurns(n int) Option {
|
||||
if n < 1 {
|
||||
n = 1
|
||||
}
|
||||
|
||||
a.maxTurns = n
|
||||
}
|
||||
}
|
||||
@@ -217,6 +219,7 @@ func WithMaxEmptyOutputRetries(n int) Option {
|
||||
if n < 0 {
|
||||
n = 0
|
||||
}
|
||||
|
||||
a.maxEmptyOutputRetries = n
|
||||
}
|
||||
}
|
||||
@@ -226,6 +229,7 @@ func WithMaxToolDepth(n int) Option {
|
||||
if n < 1 {
|
||||
n = 1
|
||||
}
|
||||
|
||||
a.maxToolDepth = n
|
||||
}
|
||||
}
|
||||
@@ -348,6 +352,7 @@ func WithMCPServers(servers ...*MCPServer) Option {
|
||||
|
||||
func WithApproval(config ApprovalConfig) Option {
|
||||
config.toolNameSet = buildToolNameSet(config.ToolNames)
|
||||
|
||||
return func(a *Agent) {
|
||||
a.approval = &config
|
||||
}
|
||||
@@ -369,6 +374,7 @@ func (a *Agent) resolveTools(ctx context.Context) ([]ToolDescriptor, map[string]
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("cannot resolve MCP tools from %q: %w", s.name, err)
|
||||
}
|
||||
|
||||
for _, t := range mcpTools {
|
||||
all = append(all, t)
|
||||
}
|
||||
@@ -380,6 +386,7 @@ func (a *Agent) resolveTools(ctx context.Context) ([]ToolDescriptor, map[string]
|
||||
if _, exists := toolMap[name]; exists {
|
||||
return nil, nil, fmt.Errorf("cannot resolve tools: duplicate tool name %q", name)
|
||||
}
|
||||
|
||||
toolMap[name] = t
|
||||
}
|
||||
|
||||
|
||||
@@ -35,8 +35,10 @@ func (m *mockProvider) ChatCompletion(_ context.Context, _ *llm.ChatCompletionRe
|
||||
if m.calls >= len(m.responses) {
|
||||
return nil, errors.New("no more mock responses")
|
||||
}
|
||||
|
||||
resp := m.responses[m.calls]
|
||||
m.calls++
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
@@ -56,6 +58,7 @@ func (s *mockChatStream) Next() bool {
|
||||
func (s *mockChatStream) Event() llm.ChatCompletionStreamEvent {
|
||||
ev := s.events[s.pos]
|
||||
s.pos++
|
||||
|
||||
return ev
|
||||
}
|
||||
|
||||
@@ -89,8 +92,10 @@ func (p *mockMultiStreamProvider) ChatCompletionStream(_ context.Context, _ *llm
|
||||
if p.calls >= len(p.streams) {
|
||||
return nil, errors.New("no more mock streams")
|
||||
}
|
||||
|
||||
s := p.streams[p.calls]
|
||||
p.calls++
|
||||
|
||||
return s, nil
|
||||
}
|
||||
|
||||
@@ -115,6 +120,7 @@ func (g *blockingGuardrail) Check(_ context.Context, messages []llm.Message) (*a
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
@@ -129,6 +135,7 @@ func (g *outputBlocker) Check(_ context.Context, message llm.Message) (*agent.Gu
|
||||
Message: "output blocked",
|
||||
}, nil
|
||||
}
|
||||
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
@@ -192,6 +199,7 @@ func (s *testSession) Load(_ context.Context, sessionID string) ([]llm.Message,
|
||||
msgs := s.messages[sessionID]
|
||||
cp := make([]llm.Message, len(msgs))
|
||||
copy(cp, msgs)
|
||||
|
||||
return cp, nil
|
||||
}
|
||||
|
||||
@@ -199,6 +207,7 @@ func (s *testSession) Save(_ context.Context, sessionID string, messages []llm.M
|
||||
cp := make([]llm.Message, len(messages))
|
||||
copy(cp, messages)
|
||||
s.messages[sessionID] = cp
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -410,6 +419,7 @@ func TestRun(t *testing.T) {
|
||||
}
|
||||
|
||||
type Params struct{}
|
||||
|
||||
noopTool := agent.FunctionTool[Params](
|
||||
"noop",
|
||||
"No-op",
|
||||
@@ -432,6 +442,7 @@ func TestRun(t *testing.T) {
|
||||
)
|
||||
|
||||
require.Error(t, err)
|
||||
|
||||
var maxTurnsErr *agent.MaxTurnsExceededError
|
||||
require.ErrorAs(t, err, &maxTurnsErr)
|
||||
assert.Equal(t, 2, maxTurnsErr.MaxTurns)
|
||||
@@ -444,6 +455,7 @@ func TestRun(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type Params struct{}
|
||||
|
||||
makeTool := func(name string) agent.Tool {
|
||||
tool := agent.FunctionTool[Params](
|
||||
name,
|
||||
@@ -452,6 +464,7 @@ func TestRun(t *testing.T) {
|
||||
return agent.ToolResult{Content: "ok"}, nil
|
||||
},
|
||||
)
|
||||
|
||||
return tool
|
||||
}
|
||||
|
||||
@@ -600,11 +613,13 @@ func TestRun(t *testing.T) {
|
||||
assert.Equal(t, "Both done.", result.FinalMessage().Text())
|
||||
|
||||
var toolMsgs []llm.Message
|
||||
|
||||
for _, m := range result.Messages {
|
||||
if m.Role == llm.RoleTool {
|
||||
toolMsgs = append(toolMsgs, m)
|
||||
}
|
||||
}
|
||||
|
||||
require.Len(t, toolMsgs, 2)
|
||||
assert.Equal(t, "tc_1", toolMsgs[0].ToolCallID)
|
||||
assert.Equal(t, "result_1", toolMsgs[0].Text())
|
||||
@@ -667,11 +682,13 @@ func TestRun(t *testing.T) {
|
||||
assert.Equal(t, "Handled both.", result.FinalMessage().Text())
|
||||
|
||||
var toolMsgs []llm.Message
|
||||
|
||||
for _, m := range result.Messages {
|
||||
if m.Role == llm.RoleTool {
|
||||
toolMsgs = append(toolMsgs, m)
|
||||
}
|
||||
}
|
||||
|
||||
require.Len(t, toolMsgs, 2)
|
||||
assert.Equal(t, "tc_ok", toolMsgs[0].ToolCallID)
|
||||
assert.Equal(t, "success_result", toolMsgs[0].Text())
|
||||
@@ -692,12 +709,14 @@ func TestRun(t *testing.T) {
|
||||
var capturedTenantID string
|
||||
|
||||
type Params struct{}
|
||||
|
||||
tool := agent.FunctionTool[Params](
|
||||
"check_tenant",
|
||||
"Check current tenant",
|
||||
func(ctx context.Context, _ Params) (agent.ToolResult, error) {
|
||||
rc := agent.RunContextFrom[*RequestContext](ctx)
|
||||
capturedTenantID = rc.TenantID
|
||||
|
||||
return agent.ToolResult{Content: "tenant: " + rc.TenantID}, nil
|
||||
},
|
||||
)
|
||||
@@ -912,11 +931,13 @@ func TestRun_Handoff(t *testing.T) {
|
||||
specialist,
|
||||
agent.WithHandoffInputFilter(func(data agent.HandoffInputData) []llm.Message {
|
||||
var filtered []llm.Message
|
||||
|
||||
for _, m := range data.NewItems {
|
||||
if m.Role == llm.RoleUser {
|
||||
filtered = append(filtered, m)
|
||||
}
|
||||
}
|
||||
|
||||
return filtered
|
||||
}),
|
||||
),
|
||||
@@ -1017,6 +1038,7 @@ func TestRun_Guardrails(t *testing.T) {
|
||||
)
|
||||
|
||||
require.Error(t, err)
|
||||
|
||||
var tripErr *agent.InputGuardrailTrippedError
|
||||
require.ErrorAs(t, err, &tripErr)
|
||||
assert.Equal(t, "blocker", tripErr.Guardrail)
|
||||
@@ -1048,6 +1070,7 @@ func TestRun_Guardrails(t *testing.T) {
|
||||
)
|
||||
|
||||
require.Error(t, err)
|
||||
|
||||
var tripErr *agent.OutputGuardrailTrippedError
|
||||
require.ErrorAs(t, err, &tripErr)
|
||||
assert.Equal(t, "output_blocker", tripErr.Guardrail)
|
||||
@@ -1064,6 +1087,7 @@ func TestRun_Hooks(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type Params struct{}
|
||||
|
||||
noopTool := agent.FunctionTool[Params](
|
||||
"noop",
|
||||
"No-op",
|
||||
@@ -1348,6 +1372,7 @@ func TestRun_ToolUseBehavior(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type Params struct{}
|
||||
|
||||
tool := agent.FunctionTool[Params](
|
||||
"compute",
|
||||
"Compute something",
|
||||
@@ -1440,6 +1465,7 @@ func TestRun_ToolUseBehavior(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type Params struct{}
|
||||
|
||||
tool := agent.FunctionTool[Params](
|
||||
"noop",
|
||||
"No-op",
|
||||
@@ -1482,6 +1508,7 @@ func TestRun_ToolUseBehavior(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type Params struct{}
|
||||
|
||||
tool := agent.FunctionTool[Params](
|
||||
"compute",
|
||||
"Compute something",
|
||||
@@ -1594,6 +1621,7 @@ func TestRun_Approval(t *testing.T) {
|
||||
)
|
||||
|
||||
require.Error(t, err)
|
||||
|
||||
var interrupted *agent.InterruptedError
|
||||
require.ErrorAs(t, err, &interrupted)
|
||||
assert.Len(t, interrupted.ToolCalls, 1)
|
||||
@@ -2107,8 +2135,10 @@ func TestRunStreamed(t *testing.T) {
|
||||
[]llm.Message{userMessage("Hi")},
|
||||
)
|
||||
|
||||
var deltas []string
|
||||
var gotComplete bool
|
||||
var (
|
||||
deltas []string
|
||||
gotComplete bool
|
||||
)
|
||||
|
||||
for ev := range sr.Events {
|
||||
switch ev.Type {
|
||||
@@ -2134,6 +2164,7 @@ func TestRunStreamed(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type Params struct{}
|
||||
|
||||
tool := agent.FunctionTool[Params](
|
||||
"noop",
|
||||
"No-op",
|
||||
@@ -2186,6 +2217,7 @@ func TestRunStreamed(t *testing.T) {
|
||||
)
|
||||
|
||||
var gotToolStart, gotToolEnd, gotComplete bool
|
||||
|
||||
for ev := range sr.Events {
|
||||
switch ev.Type {
|
||||
case agent.StreamEventToolStart:
|
||||
@@ -2238,6 +2270,7 @@ func TestRunStreamed(t *testing.T) {
|
||||
)
|
||||
|
||||
var gotComplete, gotError bool
|
||||
|
||||
for ev := range sr.Events {
|
||||
switch ev.Type {
|
||||
case agent.StreamEventComplete:
|
||||
@@ -2287,22 +2320,29 @@ func TestRunStreamed(t *testing.T) {
|
||||
)
|
||||
|
||||
var collected []agent.StreamEvent
|
||||
|
||||
done := make(chan struct{})
|
||||
|
||||
go func() {
|
||||
defer close(done)
|
||||
|
||||
for ev := range sr.Events {
|
||||
collected = append(collected, ev)
|
||||
}
|
||||
}()
|
||||
|
||||
result, err := sr.Wait()
|
||||
|
||||
<-done
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Hello world!", result.FinalMessage().Text())
|
||||
|
||||
var deltaCount int
|
||||
var gotAgentStart, gotAgentEnd, gotComplete bool
|
||||
var (
|
||||
deltaCount int
|
||||
gotAgentStart, gotAgentEnd, gotComplete bool
|
||||
)
|
||||
|
||||
for _, ev := range collected {
|
||||
switch ev.Type {
|
||||
case agent.StreamEventLLMDelta:
|
||||
@@ -2359,6 +2399,7 @@ func TestClone(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type Params struct{}
|
||||
|
||||
tool1 := agent.FunctionTool[Params](
|
||||
"t1",
|
||||
"desc",
|
||||
@@ -2474,6 +2515,7 @@ func TestGenerateSchema_EmbeddedStruct(t *testing.T) {
|
||||
ID string `json:"id" jsonschema:"unique identifier"`
|
||||
Kind string `json:"kind"`
|
||||
}
|
||||
|
||||
type Params struct {
|
||||
Base
|
||||
Name string `json:"name"`
|
||||
@@ -2653,6 +2695,7 @@ func TestRun_UnknownToolCall(t *testing.T) {
|
||||
}
|
||||
|
||||
type Params struct{}
|
||||
|
||||
tool := agent.FunctionTool[Params](
|
||||
"real_tool",
|
||||
"A real tool",
|
||||
@@ -2746,6 +2789,7 @@ func TestClone_WithApprovalConfig(t *testing.T) {
|
||||
)
|
||||
|
||||
require.Error(t, err)
|
||||
|
||||
var interrupted *agent.InterruptedError
|
||||
require.ErrorAs(t, err, &interrupted)
|
||||
assert.Len(t, interrupted.PendingApprovals, 1)
|
||||
@@ -2771,6 +2815,7 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
|
||||
var executionOrder []string
|
||||
|
||||
type Params struct{}
|
||||
|
||||
tool1 := agent.FunctionTool[Params](
|
||||
"prepare",
|
||||
"Prepare data",
|
||||
@@ -2830,6 +2875,7 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type Params struct{}
|
||||
|
||||
tool1 := agent.FunctionTool[Params](
|
||||
"prepare",
|
||||
"Prepare data",
|
||||
@@ -2892,6 +2938,7 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
|
||||
assert.Equal(t, "specialist", result.LastAgent.Name())
|
||||
|
||||
var toolMsgs []llm.Message
|
||||
|
||||
for _, m := range result.Messages {
|
||||
if m.Role == llm.RoleTool {
|
||||
toolMsgs = append(toolMsgs, m)
|
||||
@@ -2914,6 +2961,7 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type Params struct{}
|
||||
|
||||
failingTool := agent.FunctionTool[Params](
|
||||
"prepare",
|
||||
"Prepare data",
|
||||
|
||||
@@ -47,6 +47,7 @@ func agentToolDepth(ctx context.Context) int {
|
||||
if v, ok := ctx.Value(agentToolDepthKey{}).(int); ok {
|
||||
return v
|
||||
}
|
||||
|
||||
return 0
|
||||
}
|
||||
|
||||
@@ -124,6 +125,7 @@ func (t *agentTool) Execute(ctx context.Context, arguments string) (ToolResult,
|
||||
if len(preview) > 500 {
|
||||
preview = preview[:500] + "... (truncated)"
|
||||
}
|
||||
|
||||
return ToolResult{
|
||||
Content: fmt.Sprintf("Sub-agent %q returned invalid JSON. Raw output:\n%s", t.agent.name, preview),
|
||||
IsError: true,
|
||||
|
||||
@@ -273,12 +273,14 @@ func TestAgentTool_Execute(t *testing.T) {
|
||||
var captured string
|
||||
|
||||
type Params struct{}
|
||||
|
||||
tenantTool := agent.FunctionTool[Params](
|
||||
"get_tenant",
|
||||
"Get tenant",
|
||||
func(ctx context.Context, _ Params) (agent.ToolResult, error) {
|
||||
rc := agent.RunContextFrom[*AppCtx](ctx)
|
||||
captured = rc.TenantID
|
||||
|
||||
return agent.ToolResult{Content: rc.TenantID}, nil
|
||||
},
|
||||
)
|
||||
@@ -483,6 +485,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
|
||||
)
|
||||
|
||||
require.Error(t, err)
|
||||
|
||||
var interrupted *agent.InterruptedError
|
||||
require.ErrorAs(t, err, &interrupted)
|
||||
assert.Len(t, interrupted.PendingApprovals, 1)
|
||||
@@ -605,6 +608,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
|
||||
var siblingCalled bool
|
||||
|
||||
type Params struct{}
|
||||
|
||||
siblingTool := agent.FunctionTool[Params](
|
||||
"list_files",
|
||||
"List files",
|
||||
@@ -740,6 +744,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
|
||||
)
|
||||
|
||||
require.Error(t, err)
|
||||
|
||||
var interrupted *agent.InterruptedError
|
||||
require.ErrorAs(t, err, &interrupted)
|
||||
assert.Equal(t, "agent_c", interrupted.Agent.Name())
|
||||
|
||||
@@ -47,6 +47,7 @@ func buildToolNameSet(names []string) map[string]struct{} {
|
||||
for _, name := range names {
|
||||
set[name] = struct{}{}
|
||||
}
|
||||
|
||||
return set
|
||||
}
|
||||
|
||||
@@ -60,5 +61,6 @@ func (c *ApprovalConfig) requiresApproval(ctx context.Context, tc llm.ToolCall)
|
||||
}
|
||||
|
||||
_, ok := c.toolNameSet[tc.Function.Name]
|
||||
|
||||
return ok
|
||||
}
|
||||
|
||||
@@ -64,6 +64,7 @@ func TestBuildToolNameSet(t *testing.T) {
|
||||
set := buildToolNameSet([]string{"delete", "update", "create"})
|
||||
|
||||
assert.Len(t, set, 3)
|
||||
|
||||
for _, name := range []string{"delete", "update", "create"} {
|
||||
_, ok := set[name]
|
||||
assert.True(t, ok, "expected set to contain %q", name)
|
||||
@@ -165,15 +166,19 @@ func TestApprovalConfig_RequiresApproval(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type ctxKey struct{}
|
||||
|
||||
ctx := context.WithValue(context.Background(), ctxKey{}, "marker")
|
||||
|
||||
var capturedCtx context.Context
|
||||
var capturedTC llm.ToolCall
|
||||
var (
|
||||
capturedCtx context.Context
|
||||
capturedTC llm.ToolCall
|
||||
)
|
||||
|
||||
c := &ApprovalConfig{
|
||||
ShouldApprove: func(ctx context.Context, tc llm.ToolCall) bool {
|
||||
capturedCtx = ctx
|
||||
capturedTC = tc
|
||||
|
||||
return true
|
||||
},
|
||||
}
|
||||
|
||||
@@ -52,6 +52,7 @@ func (p *blockingProvider) ChatCompletion(ctx context.Context, _ *llm.ChatComple
|
||||
p.ctxAtEnd = ctx.Err()
|
||||
p.mu.Unlock()
|
||||
}
|
||||
|
||||
return p.response, nil
|
||||
}
|
||||
|
||||
@@ -214,6 +215,7 @@ func TestRun_CtxCancelGracefulSuspend(t *testing.T) {
|
||||
defer cancel()
|
||||
|
||||
done := make(chan error, 1)
|
||||
|
||||
go func() {
|
||||
_, err := ag.Run(
|
||||
ctx,
|
||||
@@ -230,6 +232,7 @@ func TestRun_CtxCancelGracefulSuspend(t *testing.T) {
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("LLM call never started")
|
||||
}
|
||||
|
||||
cancel()
|
||||
close(provider.release)
|
||||
|
||||
|
||||
@@ -49,5 +49,6 @@ func TryRunContextFrom[C any](ctx context.Context) (C, bool) {
|
||||
}
|
||||
|
||||
typed, ok := val.(C)
|
||||
|
||||
return typed, ok
|
||||
}
|
||||
|
||||
@@ -83,6 +83,7 @@ func (g *PromptInjectionGuardrail) Check(ctx context.Context, messages []llm.Mes
|
||||
"prompt injection classifier failed, allowing message through",
|
||||
log.Error(err),
|
||||
)
|
||||
|
||||
return &agent.GuardrailResult{Tripwire: false}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -32,6 +32,7 @@ func NewSystemPromptLeakGuardrail(fingerprints []string) *SystemPromptLeakGuardr
|
||||
if f == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
lowered = append(lowered, strings.ToLower(f))
|
||||
}
|
||||
|
||||
|
||||
@@ -53,6 +53,7 @@ func HandoffTo(agent *Agent, opts ...HandoffOption) *Handoff {
|
||||
for _, opt := range opts {
|
||||
opt(h)
|
||||
}
|
||||
|
||||
return h
|
||||
}
|
||||
|
||||
@@ -84,6 +85,7 @@ func (h *Handoff) toolName() string {
|
||||
if h.ToolName != "" {
|
||||
return h.ToolName
|
||||
}
|
||||
|
||||
return "transfer_to_" + sanitizeToolName(h.Agent.name)
|
||||
}
|
||||
|
||||
|
||||
@@ -156,11 +156,13 @@ func TestWithHandoffInputFilter(t *testing.T) {
|
||||
target,
|
||||
agent.WithHandoffInputFilter(func(data agent.HandoffInputData) []llm.Message {
|
||||
var filtered []llm.Message
|
||||
|
||||
for _, m := range data.NewItems {
|
||||
if m.Role == llm.RoleUser {
|
||||
filtered = append(filtered, m)
|
||||
}
|
||||
}
|
||||
|
||||
return filtered
|
||||
}),
|
||||
)
|
||||
@@ -194,6 +196,7 @@ func TestWithHandoffInputFilter(t *testing.T) {
|
||||
all := make([]llm.Message, 0, len(data.InputHistory)+len(data.NewItems))
|
||||
all = append(all, data.InputHistory...)
|
||||
all = append(all, data.NewItems...)
|
||||
|
||||
return all
|
||||
}),
|
||||
)
|
||||
|
||||
@@ -56,12 +56,15 @@ func (s *MCPServer) Name() string {
|
||||
|
||||
func (s *MCPServer) Tools(ctx context.Context) ([]Tool, error) {
|
||||
s.mu.RLock()
|
||||
|
||||
if s.toolsCached {
|
||||
cp := make([]Tool, len(s.cachedTools))
|
||||
copy(cp, s.cachedTools)
|
||||
s.mu.RUnlock()
|
||||
|
||||
return cp, nil
|
||||
}
|
||||
|
||||
s.mu.RUnlock()
|
||||
|
||||
s.mu.Lock()
|
||||
@@ -70,11 +73,14 @@ func (s *MCPServer) Tools(ctx context.Context) ([]Tool, error) {
|
||||
if s.toolsCached {
|
||||
cp := make([]Tool, len(s.cachedTools))
|
||||
copy(cp, s.cachedTools)
|
||||
|
||||
return cp, nil
|
||||
}
|
||||
|
||||
var allTools []*mcp.Tool
|
||||
var cursor string
|
||||
var (
|
||||
allTools []*mcp.Tool
|
||||
cursor string
|
||||
)
|
||||
|
||||
for {
|
||||
params := &mcp.ListToolsParams{}
|
||||
@@ -92,6 +98,7 @@ func (s *MCPServer) Tools(ctx context.Context) ([]Tool, error) {
|
||||
if result.NextCursor == "" {
|
||||
break
|
||||
}
|
||||
|
||||
cursor = result.NextCursor
|
||||
}
|
||||
|
||||
@@ -172,6 +179,7 @@ func extractMCPContent(result *mcp.CallToolResult) string {
|
||||
}
|
||||
|
||||
var parts []string
|
||||
|
||||
for _, c := range result.Content {
|
||||
if tc, ok := c.(*mcp.TextContent); ok {
|
||||
parts = append(parts, tc.Text)
|
||||
|
||||
@@ -106,6 +106,7 @@ func TestMCPServer_Tools(t *testing.T) {
|
||||
|
||||
// Mutating one slice must not affect the other.
|
||||
tools1[0] = nil
|
||||
|
||||
assert.NotNil(t, tools2[0])
|
||||
|
||||
// Underlying cache must be untouched.
|
||||
@@ -131,6 +132,7 @@ func TestMCPServer_Tools(t *testing.T) {
|
||||
s.toolsCached = true
|
||||
|
||||
const goroutines = 50
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(goroutines)
|
||||
|
||||
@@ -212,12 +214,14 @@ func TestMCPServer_ResetCache(t *testing.T) {
|
||||
s.toolsCached = true
|
||||
|
||||
const goroutines = 50
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(goroutines)
|
||||
|
||||
for range goroutines {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
|
||||
s.ResetCache()
|
||||
}()
|
||||
}
|
||||
@@ -321,6 +325,7 @@ func TestExtractMCPContent(t *testing.T) {
|
||||
"empty content returns empty",
|
||||
func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
result := &mcp.CallToolResult{Content: []mcp.Content{}}
|
||||
assert.Equal(t, "", extractMCPContent(result))
|
||||
},
|
||||
@@ -330,6 +335,7 @@ func TestExtractMCPContent(t *testing.T) {
|
||||
"single text content",
|
||||
func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
result := &mcp.CallToolResult{
|
||||
Content: []mcp.Content{
|
||||
&mcp.TextContent{Text: "hello world"},
|
||||
@@ -343,6 +349,7 @@ func TestExtractMCPContent(t *testing.T) {
|
||||
"multiple text contents joined by newline",
|
||||
func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
result := &mcp.CallToolResult{
|
||||
Content: []mcp.Content{
|
||||
&mcp.TextContent{Text: "line one"},
|
||||
@@ -357,6 +364,7 @@ func TestExtractMCPContent(t *testing.T) {
|
||||
"non-text content is skipped",
|
||||
func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
result := &mcp.CallToolResult{
|
||||
Content: []mcp.Content{
|
||||
&mcp.TextContent{Text: "text part"},
|
||||
@@ -371,6 +379,7 @@ func TestExtractMCPContent(t *testing.T) {
|
||||
"only non-text content returns empty",
|
||||
func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
result := &mcp.CallToolResult{
|
||||
Content: []mcp.Content{
|
||||
&mcp.ImageContent{Data: []byte("base64data"), MIMEType: "image/png"},
|
||||
@@ -384,6 +393,7 @@ func TestExtractMCPContent(t *testing.T) {
|
||||
"falls back to structured content when no text content",
|
||||
func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
result := &mcp.CallToolResult{
|
||||
StructuredContent: map[string]any{
|
||||
"status": "ok",
|
||||
@@ -400,6 +410,7 @@ func TestExtractMCPContent(t *testing.T) {
|
||||
"text content takes precedence over structured content",
|
||||
func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
result := &mcp.CallToolResult{
|
||||
Content: []mcp.Content{
|
||||
&mcp.TextContent{Text: "text wins"},
|
||||
@@ -414,6 +425,7 @@ func TestExtractMCPContent(t *testing.T) {
|
||||
"structured content used when content has only non-text",
|
||||
func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
result := &mcp.CallToolResult{
|
||||
Content: []mcp.Content{
|
||||
&mcp.ImageContent{Data: []byte("img"), MIMEType: "image/png"},
|
||||
|
||||
@@ -75,6 +75,7 @@ func TestOutputType_responseFormat(t *testing.T) {
|
||||
|
||||
ot, err := NewOutputType[Verdict]("verdict")
|
||||
require.NoError(t, err)
|
||||
|
||||
rf := ot.responseFormat()
|
||||
|
||||
require.NotNil(t, rf)
|
||||
@@ -96,6 +97,7 @@ func TestOutputType_responseFormat_SchemaMatchesOutputType(t *testing.T) {
|
||||
|
||||
ot, err := NewOutputType[Analysis]("analysis")
|
||||
require.NoError(t, err)
|
||||
|
||||
rf := ot.responseFormat()
|
||||
|
||||
var schema map[string]any
|
||||
|
||||
@@ -37,13 +37,16 @@ func Restore(
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot load checkpoint: %w", err)
|
||||
}
|
||||
|
||||
if cp == nil {
|
||||
return nil, fmt.Errorf("cannot restore: no checkpoint for run %s", runID)
|
||||
}
|
||||
|
||||
agent, err := registry.Agent(cp.AgentName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot resolve agent %q: %w", cp.AgentName, err)
|
||||
}
|
||||
|
||||
agent = applyCheckpointConfig(agent, cp.Config)
|
||||
|
||||
return restoreCheckpoint(ctx, agent, cp, store, runID, registry)
|
||||
@@ -58,6 +61,7 @@ func applyCheckpointConfig(agent *Agent, cfg AgentConfig) *Agent {
|
||||
if cfg.MaxTurns <= 0 {
|
||||
return agent
|
||||
}
|
||||
|
||||
return agent.Clone(WithMaxTurns(cfg.MaxTurns))
|
||||
}
|
||||
|
||||
@@ -154,13 +158,17 @@ func restoreNestedSuspended(
|
||||
}
|
||||
|
||||
entries := make([]nestedRestoreEntry, len(cp.AllToolCalls))
|
||||
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for i, tc := range cp.AllToolCalls {
|
||||
entries[i].toolCall = tc
|
||||
|
||||
result, ok := completedByID[tc.ID]
|
||||
if ok {
|
||||
entries[i].result = result
|
||||
entries[i].completed = true
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -169,6 +177,7 @@ func restoreNestedSuspended(
|
||||
entries[i].err = fmt.Errorf("cannot restore nested tool call %q: missing inner checkpoint", tc.ID)
|
||||
continue
|
||||
}
|
||||
|
||||
entries[i].originalCheckpoint = innerCP
|
||||
|
||||
innerAgent, err := registry.Agent(innerCP.AgentName)
|
||||
@@ -176,9 +185,11 @@ func restoreNestedSuspended(
|
||||
entries[i].err = fmt.Errorf("cannot resolve inner agent %q: %w", innerCP.AgentName, err)
|
||||
continue
|
||||
}
|
||||
|
||||
innerAgent = applyCheckpointConfig(innerAgent, innerCP.Config)
|
||||
|
||||
wg.Add(1)
|
||||
|
||||
go func(i int, tc llm.ToolCall, innerAgent *Agent, innerCP *Checkpoint) {
|
||||
defer wg.Done()
|
||||
|
||||
@@ -189,10 +200,14 @@ func restoreNestedSuspended(
|
||||
entries[i].err = fmt.Errorf("cannot restore nested tool call %q: missing suspension checkpoint", tc.ID)
|
||||
return
|
||||
}
|
||||
|
||||
entries[i].suspendedCheckpoint = se.Checkpoint
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
entries[i].err = fmt.Errorf("cannot restore nested tool call %q: %w", tc.ID, err)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
@@ -200,6 +215,7 @@ func restoreNestedSuspended(
|
||||
entries[i].completed = true
|
||||
}(i, tc, innerAgent, innerCP)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
messages := make([]llm.Message, len(cp.Messages))
|
||||
@@ -207,16 +223,20 @@ func restoreNestedSuspended(
|
||||
|
||||
completedCalls := make([]CompletedCall, 0, len(cp.AllToolCalls))
|
||||
remainingInner := make(map[string]*Checkpoint)
|
||||
|
||||
var restoreErr error
|
||||
|
||||
for _, entry := range entries {
|
||||
switch {
|
||||
case entry.err != nil:
|
||||
if entry.originalCheckpoint != nil {
|
||||
remainingInner[entry.toolCall.ID] = entry.originalCheckpoint
|
||||
}
|
||||
|
||||
if restoreErr == nil {
|
||||
restoreErr = entry.err
|
||||
}
|
||||
|
||||
continue
|
||||
|
||||
case entry.suspendedCheckpoint != nil:
|
||||
@@ -227,6 +247,7 @@ func restoreNestedSuspended(
|
||||
if restoreErr == nil {
|
||||
restoreErr = fmt.Errorf("cannot restore nested tool call %q: no result", entry.toolCall.ID)
|
||||
}
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -250,13 +271,16 @@ func restoreNestedSuspended(
|
||||
saveProgress := func() (*Checkpoint, error) {
|
||||
next := *cp
|
||||
next.InnerCheckpoints = remainingInner
|
||||
|
||||
next.CompletedCalls = completedCalls
|
||||
if store != nil && runID != "" {
|
||||
if err := store.Save(saveCtx, runID, &next); err != nil {
|
||||
return nil, fmt.Errorf("cannot save nested restore progress: %w", err)
|
||||
}
|
||||
|
||||
emitHook(agent, func(h RunHooks) { h.OnRunSnapshot(saveCtx, agent, &next) })
|
||||
}
|
||||
|
||||
return &next, nil
|
||||
}
|
||||
|
||||
@@ -264,6 +288,7 @@ func restoreNestedSuspended(
|
||||
if _, err := saveProgress(); err != nil {
|
||||
return nil, errors.Join(restoreErr, err)
|
||||
}
|
||||
|
||||
return nil, restoreErr
|
||||
}
|
||||
|
||||
@@ -272,6 +297,7 @@ func restoreNestedSuspended(
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return nil, &SuspendedError{RunID: runID, Checkpoint: next}
|
||||
}
|
||||
|
||||
@@ -301,11 +327,13 @@ func restoreAwaitingApproval(
|
||||
if len(cp.InnerCheckpoints) > 1 {
|
||||
return nil, fmt.Errorf("cannot restore approval checkpoint: expected one inner checkpoint, got %d", len(cp.InnerCheckpoints))
|
||||
}
|
||||
|
||||
for toolCallID, innerCP := range cp.InnerCheckpoints {
|
||||
innerAgent, err := registry.Agent(innerCP.AgentName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot resolve inner agent %q: %w", innerCP.AgentName, err)
|
||||
}
|
||||
|
||||
innerAgent = applyCheckpointConfig(innerAgent, innerCP.Config)
|
||||
|
||||
innerIE := &InterruptedError{
|
||||
@@ -334,6 +362,7 @@ func restoreAwaitingApproval(
|
||||
completedCalls: cp.CompletedCalls,
|
||||
innerInterrupt: innerIE,
|
||||
}
|
||||
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
@@ -43,6 +43,7 @@ func (s *memoryCheckpointer) Save(_ context.Context, runID string, cp *agent.Che
|
||||
|
||||
clone := *cp
|
||||
s.checkpoints[runID] = &clone
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -56,6 +57,7 @@ func (s *memoryCheckpointer) Load(_ context.Context, runID string) (*agent.Check
|
||||
}
|
||||
|
||||
clone := *cp
|
||||
|
||||
return &clone, nil
|
||||
}
|
||||
|
||||
@@ -68,6 +70,7 @@ func (r *simpleRegistry) Agent(name string) (*agent.Agent, error) {
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("agent %q not found", name)
|
||||
}
|
||||
|
||||
return a, nil
|
||||
}
|
||||
|
||||
@@ -220,6 +223,7 @@ func TestRestore(t *testing.T) {
|
||||
)
|
||||
|
||||
require.Error(t, err)
|
||||
|
||||
var interrupted *agent.InterruptedError
|
||||
require.ErrorAs(t, err, &interrupted)
|
||||
assert.Len(t, interrupted.PendingApprovals, 1)
|
||||
|
||||
@@ -29,5 +29,6 @@ func (r *Result) FinalMessage() llm.Message {
|
||||
if len(r.Messages) == 0 {
|
||||
return llm.Message{}
|
||||
}
|
||||
|
||||
return r.Messages[len(r.Messages)-1]
|
||||
}
|
||||
|
||||
@@ -104,14 +104,17 @@ func blockingCallLLM(ctx context.Context, agent *Agent, req *llm.ChatCompletionR
|
||||
if sErr != nil {
|
||||
return nil, err // return the original error
|
||||
}
|
||||
|
||||
defer func() { _ = stream.Close() }()
|
||||
|
||||
acc := llm.NewStreamAccumulator(stream)
|
||||
for acc.Next() {
|
||||
}
|
||||
|
||||
if sErr := acc.Err(); sErr != nil {
|
||||
return nil, sErr
|
||||
}
|
||||
|
||||
return acc.Response(), nil
|
||||
}
|
||||
|
||||
@@ -143,6 +146,7 @@ func (s *loopState) resolveAgentTools(ctx context.Context) error {
|
||||
|
||||
s.toolMap = toolMap
|
||||
s.toolDefs = toolDefs
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -218,6 +222,7 @@ func (s *loopState) finishRun(ctx context.Context, result *Result, err error) (*
|
||||
)
|
||||
|
||||
s.opts.onEvent(ctx, StreamEvent{Type: StreamEventComplete, Agent: s.agent, Result: result})
|
||||
|
||||
return result, err
|
||||
}
|
||||
|
||||
@@ -385,6 +390,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
|
||||
if s.opts.checkpointer != nil {
|
||||
if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil {
|
||||
s.logger.ErrorCtx(ctx, "cannot save suspension checkpoint", log.Error(saveErr))
|
||||
|
||||
se.Checkpoint = cp
|
||||
} else {
|
||||
emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, cp) })
|
||||
@@ -412,6 +418,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
|
||||
if s.toolUsedInRun && s.agent.resetToolChoice && toolChoice != nil {
|
||||
toolChoice = nil
|
||||
}
|
||||
|
||||
if !exploring && structuredFormat != nil && len(s.toolDefs) > 0 {
|
||||
// On the synthesis turn, forbid further tool calls so the
|
||||
// model is forced to convert what it has into JSON.
|
||||
@@ -478,9 +485,11 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
|
||||
// conclusions during synthesis.
|
||||
if exploring && s.turns < s.agent.maxTurns {
|
||||
exploring = false
|
||||
|
||||
if resp.Message.Text() == "" {
|
||||
s.messages = s.messages[:len(s.messages)-1]
|
||||
}
|
||||
|
||||
s.messages = append(
|
||||
s.messages,
|
||||
llm.Message{
|
||||
@@ -494,6 +503,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
|
||||
log.Int("turn", s.turns),
|
||||
log.Int("output_tokens", resp.Usage.OutputTokens),
|
||||
)
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -514,8 +524,10 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
|
||||
log.Int("retry", emptyOutputRetries),
|
||||
log.Int("output_tokens", resp.Usage.OutputTokens),
|
||||
)
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
if err := runOutputGuardrails(ctx, s.agent, resp.Message); err != nil {
|
||||
return s.finishRun(ctx, nil, err)
|
||||
}
|
||||
@@ -530,6 +542,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
|
||||
emitAgentHook(s.agent, func(h AgentHooks) { h.OnEnd(ctx, s.agent, resp.Message.Text()) })
|
||||
|
||||
opts.onEvent(ctx, StreamEvent{Type: StreamEventAgentEnd, Agent: s.agent})
|
||||
|
||||
return s.finishRun(ctx, result, nil)
|
||||
|
||||
case llm.FinishReasonToolCalls:
|
||||
@@ -561,6 +574,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
|
||||
outerCP.InnerCheckpoints = se.Checkpoint.InnerCheckpoints
|
||||
outerCP.CompletedCalls = se.Checkpoint.CompletedCalls
|
||||
}
|
||||
|
||||
if s.opts.checkpointer != nil {
|
||||
if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, outerCP); saveErr != nil {
|
||||
s.logger.ErrorCtx(ctx, "cannot save checkpoint", log.Error(saveErr))
|
||||
@@ -568,6 +582,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
|
||||
emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, outerCP) })
|
||||
}
|
||||
}
|
||||
|
||||
return s.finishRun(ctx, nil, &SuspendedError{RunID: s.opts.runID, Checkpoint: outerCP})
|
||||
}
|
||||
|
||||
@@ -584,6 +599,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
|
||||
if s.opts.checkpointer != nil {
|
||||
cp := s.buildCheckpoint(AgentStatusAwaitingApproval)
|
||||
cp.PendingToolCalls = nae.allToolCalls
|
||||
|
||||
cp.PendingApprovals = nae.pendingApprovals
|
||||
if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil {
|
||||
s.logger.ErrorCtx(ctx, "cannot save approval checkpoint", log.Error(saveErr))
|
||||
@@ -623,6 +639,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
|
||||
cp.PendingApprovals = nie.inner.PendingApprovals
|
||||
cp.AllToolCalls = nie.allToolCalls
|
||||
cp.CompletedCalls = nie.completedCalls
|
||||
|
||||
cp.InnerCheckpoints = map[string]*Checkpoint{
|
||||
nie.toolCallID: {
|
||||
Status: AgentStatusAwaitingApproval,
|
||||
@@ -692,6 +709,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
|
||||
emitAgentHook(s.agent, func(h AgentHooks) { h.OnEnd(ctx, s.agent, finalOutput) })
|
||||
|
||||
opts.onEvent(ctx, StreamEvent{Type: StreamEventAgentEnd, Agent: s.agent})
|
||||
|
||||
return s.finishRun(ctx, result, nil)
|
||||
}
|
||||
|
||||
@@ -748,6 +766,7 @@ func callLLMWithHooks(
|
||||
if err != nil {
|
||||
emitHook(agent, func(h RunHooks) { h.OnLLMEnd(ctx, agent, nil, err) })
|
||||
emitAgentHook(agent, func(h AgentHooks) { h.OnLLMEnd(ctx, agent, nil, err) })
|
||||
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -790,6 +809,7 @@ func executeToolCalls(
|
||||
if !ok {
|
||||
return nil, nil, nil, fmt.Errorf("cannot dispatch tool call: unknown tool %q", tc.Function.Name)
|
||||
}
|
||||
|
||||
descriptors[i] = desc
|
||||
if _, isHandoff := desc.(*handoffToolAdapter); isHandoff && handoffIdx == -1 {
|
||||
handoffIdx = i
|
||||
@@ -806,6 +826,7 @@ func executeToolCalls(
|
||||
}
|
||||
|
||||
results, msgs, err := executeParallel(ctx, tracer, agent, toolCalls, tools, onEvent, logger)
|
||||
|
||||
return nil, results, msgs, err
|
||||
}
|
||||
|
||||
@@ -844,6 +865,7 @@ func executeWithHandoff(
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
return nil, nil, msgs, &nestedInterruptionError{
|
||||
inner: ie,
|
||||
toolCallID: toolCalls[i].ID,
|
||||
@@ -851,6 +873,7 @@ func executeWithHandoff(
|
||||
completedCalls: completed,
|
||||
}
|
||||
}
|
||||
|
||||
return nil, nil, msgs, err
|
||||
}
|
||||
|
||||
@@ -915,9 +938,11 @@ func executeParallel(
|
||||
logger *log.Logger,
|
||||
) ([]ToolCallResult, []llm.Message, error) {
|
||||
entries := make([]parallelToolEntry, len(toolCalls))
|
||||
|
||||
var wg sync.WaitGroup
|
||||
|
||||
wg.Add(len(toolCalls))
|
||||
|
||||
for i := range toolCalls {
|
||||
go func(idx int, tc llm.ToolCall, tool Tool) {
|
||||
defer wg.Done()
|
||||
@@ -927,6 +952,7 @@ func executeParallel(
|
||||
entries[idx] = parallelToolEntry{err: err}
|
||||
return
|
||||
}
|
||||
|
||||
entries[idx] = parallelToolEntry{result: tr}
|
||||
}(i, toolCalls[i], tools[i])
|
||||
}
|
||||
@@ -940,10 +966,12 @@ func executeParallel(
|
||||
}
|
||||
|
||||
var completed []CompletedCall
|
||||
|
||||
for j, other := range entries {
|
||||
if j == i {
|
||||
continue
|
||||
}
|
||||
|
||||
if other.err != nil {
|
||||
completed = append(
|
||||
completed,
|
||||
@@ -955,8 +983,10 @@ func executeParallel(
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
completed = append(
|
||||
completed,
|
||||
CompletedCall{
|
||||
@@ -965,6 +995,7 @@ func executeParallel(
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
return nil, nil, &nestedInterruptionError{
|
||||
inner: ie,
|
||||
toolCallID: toolCalls[i].ID,
|
||||
@@ -978,15 +1009,18 @@ func executeParallel(
|
||||
if entry.err == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
se, ok := errors.AsType[*SuspendedError](entry.err)
|
||||
if ok && se.Checkpoint != nil {
|
||||
innerCheckpoints := make(map[string]*Checkpoint)
|
||||
|
||||
var completed []CompletedCall
|
||||
|
||||
for j, other := range entries {
|
||||
if j == i {
|
||||
continue
|
||||
}
|
||||
|
||||
if other.err == nil {
|
||||
completed = append(
|
||||
completed,
|
||||
@@ -995,8 +1029,10 @@ func executeParallel(
|
||||
Result: other.result,
|
||||
},
|
||||
)
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
otherSE, ok := errors.AsType[*SuspendedError](other.err)
|
||||
if ok && otherSE.Checkpoint != nil {
|
||||
innerCheckpoints[toolCalls[j].ID] = otherSE.Checkpoint
|
||||
@@ -1024,6 +1060,7 @@ func executeParallel(
|
||||
CompletedCalls: completed,
|
||||
},
|
||||
}
|
||||
|
||||
return nil, nil, outerSE
|
||||
}
|
||||
}
|
||||
@@ -1060,6 +1097,7 @@ func executeParallel(
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -1169,6 +1207,7 @@ func executeSingleTool(
|
||||
if len(content) > 200 {
|
||||
content = content[:200] + "... (truncated)"
|
||||
}
|
||||
|
||||
logger.WarnCtx(
|
||||
ctx,
|
||||
"tool returned error",
|
||||
@@ -1192,6 +1231,7 @@ func checkApproval(ctx context.Context, a *Agent, toolCalls []llm.ToolCall) erro
|
||||
}
|
||||
|
||||
var pending []llm.ToolCall
|
||||
|
||||
for _, tc := range toolCalls {
|
||||
if a.approval.requiresApproval(ctx, tc) {
|
||||
pending = append(pending, tc)
|
||||
@@ -1363,6 +1403,7 @@ func resumeWithOpts(ctx context.Context, interrupted *InterruptedError, input Re
|
||||
)
|
||||
|
||||
handoffTarget = ht.handoff
|
||||
|
||||
break
|
||||
}
|
||||
|
||||
@@ -1464,6 +1505,7 @@ func resumeNested(ctx context.Context, interrupted *InterruptedError, input Resu
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("cannot resume nested agent: %w", err)
|
||||
}
|
||||
|
||||
@@ -1534,8 +1576,10 @@ func resolveStructuredFormat(a *Agent) *llm.ResponseFormat {
|
||||
if a.responseFormat != nil {
|
||||
return a.responseFormat
|
||||
}
|
||||
|
||||
if a.outputType != nil {
|
||||
return a.outputType.responseFormat()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -45,6 +45,7 @@ func mustJSONSchemaFor[T any]() json.RawMessage {
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
return schema
|
||||
}
|
||||
|
||||
@@ -63,6 +64,7 @@ func stripNullTypes(s *jsonschema.Schema) {
|
||||
filtered = append(filtered, t)
|
||||
}
|
||||
}
|
||||
|
||||
if len(filtered) == 1 {
|
||||
s.Type = filtered[0]
|
||||
s.Types = nil
|
||||
|
||||
@@ -129,6 +129,7 @@ func TestGenerateSchema_NestedPointerStruct(t *testing.T) {
|
||||
type Inner struct {
|
||||
Value *string `json:"value"`
|
||||
}
|
||||
|
||||
type Params struct {
|
||||
Inner *Inner `json:"inner"`
|
||||
}
|
||||
@@ -251,12 +252,15 @@ func TestGenerateSchema_DeeplyNestedStructure(t *testing.T) {
|
||||
type Level3 struct {
|
||||
Value *int `json:"value"`
|
||||
}
|
||||
|
||||
type Level2 struct {
|
||||
Items []Level3 `json:"items"`
|
||||
}
|
||||
|
||||
type Level1 struct {
|
||||
Child *Level2 `json:"child"`
|
||||
}
|
||||
|
||||
type Params struct {
|
||||
Root Level1 `json:"root"`
|
||||
}
|
||||
@@ -299,6 +303,7 @@ func TestGenerateSchema_SliceOfStructs(t *testing.T) {
|
||||
Name string `json:"name"`
|
||||
Count *int `json:"count,omitempty"`
|
||||
}
|
||||
|
||||
type Params struct {
|
||||
Items []Item `json:"items"`
|
||||
}
|
||||
|
||||
@@ -61,6 +61,7 @@ func (s *memorySession) Save(_ context.Context, sessionID string, messages []llm
|
||||
}
|
||||
|
||||
s.sessions[sessionID] = cp
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -45,6 +45,7 @@ You can transfer the conversation to a more specialized agent when appropriate:
|
||||
|
||||
func buildSystemPrompt(data systemPromptData) string {
|
||||
var buf bytes.Buffer
|
||||
|
||||
_ = systemPromptTmpl.Execute(&buf, data)
|
||||
|
||||
return buf.String()
|
||||
|
||||
@@ -50,6 +50,7 @@ func ResultJSON(v any) ToolResult {
|
||||
IsError: true,
|
||||
}
|
||||
}
|
||||
|
||||
return ToolResult{Content: string(data)}
|
||||
}
|
||||
|
||||
@@ -128,6 +129,7 @@ func (t *functionTool[P]) Execute(ctx context.Context, arguments string) (ToolRe
|
||||
}
|
||||
|
||||
var missing []string
|
||||
|
||||
for _, f := range t.requiredFields {
|
||||
if _, ok := fields[f]; !ok {
|
||||
missing = append(missing, f)
|
||||
|
||||
@@ -193,6 +193,7 @@ func TestFunctionTool_Execute(t *testing.T) {
|
||||
}
|
||||
|
||||
var received string
|
||||
|
||||
tool := agent.FunctionTool(
|
||||
"weather",
|
||||
"Get weather",
|
||||
@@ -257,6 +258,7 @@ func TestFunctionTool_Execute(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type ctxKey struct{}
|
||||
|
||||
type Params struct{}
|
||||
|
||||
tool := agent.FunctionTool(
|
||||
|
||||
@@ -42,6 +42,7 @@ func StopOnFirstTool() ToolUseBehavior {
|
||||
if len(results) == 0 {
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
return results[0].Result.Content, true, nil
|
||||
}
|
||||
}
|
||||
@@ -53,12 +54,14 @@ func StopAtTools(names ...string) ToolUseBehavior {
|
||||
for _, n := range names {
|
||||
stopSet[n] = struct{}{}
|
||||
}
|
||||
|
||||
return func(_ context.Context, results []ToolCallResult) (string, bool, error) {
|
||||
for _, r := range results {
|
||||
if _, ok := stopSet[r.ToolName]; ok {
|
||||
return r.Result.Content, true, nil
|
||||
}
|
||||
}
|
||||
|
||||
return "", false, nil
|
||||
}
|
||||
}
|
||||
|
||||
@@ -123,6 +123,7 @@ func (b *Browser) checkAlive() *agent.ToolResult {
|
||||
IsError: true,
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -72,6 +72,7 @@ func DownloadPDFTool() agent.Tool {
|
||||
ErrorDetail: fmt.Sprintf("cannot download PDF: %s", err),
|
||||
}), nil
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
@@ -95,6 +96,7 @@ func DownloadPDFTool() agent.Tool {
|
||||
ErrorDetail: fmt.Sprintf("cannot create temp dir: %s", err),
|
||||
}), nil
|
||||
}
|
||||
|
||||
defer func() { _ = os.RemoveAll(tmpDir) }()
|
||||
|
||||
tmpFile := filepath.Join(tmpDir, "input.pdf")
|
||||
@@ -106,6 +108,7 @@ func DownloadPDFTool() agent.Tool {
|
||||
|
||||
// Get page count.
|
||||
conf := model.NewDefaultConfiguration()
|
||||
|
||||
pageCount, err := api.PageCountFile(tmpFile)
|
||||
if err != nil {
|
||||
return agent.ResultJSON(downloadPDFResult{
|
||||
@@ -130,15 +133,18 @@ func DownloadPDFTool() agent.Tool {
|
||||
|
||||
// Read all extracted content files.
|
||||
var sb strings.Builder
|
||||
|
||||
entries, _ := os.ReadDir(outDir)
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
|
||||
content, err := os.ReadFile(filepath.Join(outDir, entry.Name()))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
sb.Write(content)
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
|
||||
@@ -69,6 +69,7 @@ func FetchRobotsTxtTool() agent.Tool {
|
||||
ErrorDetail: fmt.Sprintf("cannot fetch robots.txt: %s", err),
|
||||
}), nil
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
@@ -79,6 +80,7 @@ func FetchRobotsTxtTool() agent.Tool {
|
||||
}
|
||||
|
||||
var result robotsResult
|
||||
|
||||
result.Found = true
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
|
||||
@@ -73,6 +73,7 @@ func FetchSitemapTool() agent.Tool {
|
||||
ErrorDetail: fmt.Sprintf("cannot fetch sitemap: %s", err),
|
||||
}), nil
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
@@ -92,7 +93,9 @@ func FetchSitemapTool() agent.Tool {
|
||||
ErrorDetail: fmt.Sprintf("cannot decompress gzipped sitemap: %s", err),
|
||||
}), nil
|
||||
}
|
||||
|
||||
defer func() { _ = gz.Close() }()
|
||||
|
||||
reader = gz
|
||||
}
|
||||
|
||||
@@ -125,6 +128,7 @@ func FetchSitemapTool() agent.Tool {
|
||||
|
||||
func parseSitemapXML(r io.Reader) ([]string, error) {
|
||||
var urls []string
|
||||
|
||||
decoder := xml.NewDecoder(r)
|
||||
|
||||
for {
|
||||
@@ -132,6 +136,7 @@ func parseSitemapXML(r io.Reader) ([]string, error) {
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return urls, err
|
||||
}
|
||||
|
||||
@@ -119,7 +119,9 @@ func NewPinnedTransport() *http.Transport {
|
||||
|
||||
// Dial the first validated IP directly to prevent DNS rebinding.
|
||||
pinnedAddr := net.JoinHostPort(ips[0].IP.String(), port)
|
||||
|
||||
var d net.Dialer
|
||||
|
||||
return d.DialContext(ctx, network, pinnedAddr)
|
||||
},
|
||||
}
|
||||
|
||||
@@ -52,6 +52,7 @@ func DiffDocumentsTool() agent.Tool {
|
||||
if labelA == "" {
|
||||
labelA = "document_a"
|
||||
}
|
||||
|
||||
labelB := p.LabelB
|
||||
if labelB == "" {
|
||||
labelB = "document_b"
|
||||
@@ -80,6 +81,7 @@ func DiffDocumentsTool() agent.Tool {
|
||||
if len(output) > maxDiffOutput {
|
||||
output = output[:maxDiffOutput] + "\n[... diff truncated]"
|
||||
}
|
||||
|
||||
result.UnifiedDiff = output
|
||||
}
|
||||
|
||||
@@ -114,6 +116,7 @@ func computeDiff(linesA, linesB []string, labelA, labelB string) diffOutput {
|
||||
for i := range dp {
|
||||
dp[i] = make([]int, n+1)
|
||||
}
|
||||
|
||||
for i := m - 1; i >= 0; i-- {
|
||||
for j := n - 1; j >= 0; j-- {
|
||||
if linesA[i] == linesB[j] {
|
||||
@@ -131,6 +134,7 @@ func computeDiff(linesA, linesB []string, labelA, labelB string) diffOutput {
|
||||
fmt.Fprintf(&sb, "--- %s\n+++ %s\n", labelA, labelB)
|
||||
|
||||
var added, removed int
|
||||
|
||||
i, j := 0, 0
|
||||
for i < m || j < n {
|
||||
if i < m && j < n && linesA[i] == linesB[j] {
|
||||
@@ -139,10 +143,12 @@ func computeDiff(linesA, linesB []string, labelA, labelB string) diffOutput {
|
||||
j++
|
||||
} else if j < n && (i >= m || dp[i][j+1] >= dp[i+1][j]) {
|
||||
sb.WriteString("+ " + linesB[j] + "\n")
|
||||
|
||||
added++
|
||||
j++
|
||||
} else if i < m {
|
||||
sb.WriteString("- " + linesA[i] + "\n")
|
||||
|
||||
removed++
|
||||
i++
|
||||
}
|
||||
|
||||
@@ -72,6 +72,7 @@ func FirecrawlSearchTool(apiKey string) agent.Tool {
|
||||
if maxResults <= 0 {
|
||||
maxResults = 5
|
||||
}
|
||||
|
||||
if maxResults > 10 {
|
||||
maxResults = 10
|
||||
}
|
||||
@@ -111,6 +112,7 @@ func firecrawlSearch(
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
|
||||
@@ -118,6 +120,7 @@ func firecrawlSearch(
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot execute search request: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
|
||||
@@ -91,6 +91,7 @@ func CheckGovernmentDBTool(apiKey string) agent.Tool {
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
for _, e := range entries {
|
||||
*s.target = append(
|
||||
*s.target,
|
||||
|
||||
@@ -28,6 +28,7 @@ type userAgentTransport struct {
|
||||
func (t *userAgentTransport) RoundTrip(r *http.Request) (*http.Response, error) {
|
||||
r2 := r.Clone(r.Context())
|
||||
r2.Header.Set("User-Agent", "Probo-Agent/1.0")
|
||||
|
||||
return t.next.RoundTrip(r2)
|
||||
}
|
||||
|
||||
@@ -35,5 +36,6 @@ func newHTTPClient() *http.Client {
|
||||
client := httpclient.DefaultPooledClient()
|
||||
client.Timeout = 15 * time.Second
|
||||
client.Transport = &userAgentTransport{next: client.Transport}
|
||||
|
||||
return client
|
||||
}
|
||||
|
||||
@@ -66,6 +66,7 @@ func CheckWaybackTool() agent.Tool {
|
||||
|
||||
// Check availability.
|
||||
availURL := "https://archive.org/wayback/available?url=" + url.QueryEscape(p.URL)
|
||||
|
||||
body, err := httpGet(ctx, client, availURL)
|
||||
if err != nil {
|
||||
result.ErrorDetail = fmt.Sprintf("cannot check Wayback Machine availability: %s", err)
|
||||
@@ -118,6 +119,7 @@ func httpGet(ctx context.Context, client *http.Client, rawURL string) ([]byte, e
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
|
||||
@@ -101,6 +101,7 @@ func CheckCORSTool() agent.Tool {
|
||||
ErrorDetail: fmt.Sprintf("cannot fetch %s: %s", p.URL, err),
|
||||
}), nil
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
allowOrigin := resp.Header.Get("Access-Control-Allow-Origin")
|
||||
|
||||
@@ -92,6 +92,7 @@ func AnalyzeCSPTool() agent.Tool {
|
||||
ErrorDetail: fmt.Sprintf("cannot fetch %s: %s", p.URL, err),
|
||||
}), nil
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
raw := resp.Header.Get("Content-Security-Policy")
|
||||
@@ -111,6 +112,7 @@ func AnalyzeCSPTool() agent.Tool {
|
||||
directives := parseCSPDirectives(raw)
|
||||
|
||||
var hasUnsafeEval, hasUnsafeInline, hasWildcard bool
|
||||
|
||||
for _, d := range directives {
|
||||
for _, v := range d.Values {
|
||||
switch v {
|
||||
|
||||
@@ -46,6 +46,7 @@ func parseDMARCTag(record, tag string) string {
|
||||
return after
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
@@ -60,6 +61,7 @@ func CheckDMARCTool() agent.Tool {
|
||||
}
|
||||
|
||||
client := dns.NewClient()
|
||||
|
||||
answers, err := queryDNS(
|
||||
ctx,
|
||||
client,
|
||||
|
||||
@@ -53,10 +53,14 @@ func CheckDNSRecordsTool() agent.Tool {
|
||||
|
||||
hdr := dns.Header{Name: fqdn, Class: dns.ClassINET}
|
||||
client := dns.NewClient()
|
||||
var result dnsRecordsResult
|
||||
var errs []string
|
||||
|
||||
var (
|
||||
result dnsRecordsResult
|
||||
errs []string
|
||||
)
|
||||
|
||||
// A records.
|
||||
|
||||
if answers, err := queryDNS(ctx, client, &dns.A{Hdr: hdr}); err != nil {
|
||||
errs = append(errs, fmt.Sprintf("A query failed: %s", err))
|
||||
} else {
|
||||
@@ -148,12 +152,14 @@ func queryDNS(ctx context.Context, client *dns.Client, question dns.RR, opts ...
|
||||
for _, opt := range opts {
|
||||
opt(&msg.MsgHeader)
|
||||
}
|
||||
|
||||
msg.Question = []dns.RR{question}
|
||||
|
||||
resp, _, err := client.Exchange(ctx, msg, "udp", defaultResolverAddr)
|
||||
if err == nil && resp.Truncated {
|
||||
resp, _, err = client.Exchange(ctx, msg, "tcp", defaultResolverAddr)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -48,6 +48,7 @@ func CheckDNSSECTool() agent.Tool {
|
||||
}
|
||||
|
||||
client := dns.NewClient()
|
||||
|
||||
answers, err := queryDNS(
|
||||
ctx,
|
||||
client,
|
||||
@@ -66,8 +67,11 @@ func CheckDNSSECTool() agent.Tool {
|
||||
}), nil
|
||||
}
|
||||
|
||||
var keyCount int
|
||||
var keyDetails []string
|
||||
var (
|
||||
keyCount int
|
||||
keyDetails []string
|
||||
)
|
||||
|
||||
for _, answer := range answers {
|
||||
if key, ok := answer.(*dns.DNSKEY); ok {
|
||||
keyCount++
|
||||
@@ -76,6 +80,7 @@ func CheckDNSSECTool() agent.Tool {
|
||||
if key.Flags&0x0001 != 0 {
|
||||
flags = "KSK"
|
||||
}
|
||||
|
||||
keyDetails = append(
|
||||
keyDetails,
|
||||
fmt.Sprintf("%s (algorithm=%d, flags=%d)", flags, key.Algorithm, key.Flags),
|
||||
|
||||
@@ -52,6 +52,7 @@ type (
|
||||
|
||||
func checkHeader(h http.Header, name string) headerCheck {
|
||||
v := h.Get(name)
|
||||
|
||||
return headerCheck{
|
||||
Present: v != "",
|
||||
Value: v,
|
||||
@@ -92,6 +93,7 @@ func CheckSecurityHeadersTool() agent.Tool {
|
||||
|
||||
// First check the HTTP version to detect HTTP→HTTPS redirect.
|
||||
redirectsToHTTPS := false
|
||||
|
||||
httpURL := p.URL
|
||||
if after, ok := strings.CutPrefix(httpURL, "https://"); ok {
|
||||
httpURL = "http://" + after
|
||||
@@ -118,18 +120,21 @@ func CheckSecurityHeadersTool() agent.Tool {
|
||||
}
|
||||
|
||||
followClient := &http.Client{Timeout: 10 * time.Second}
|
||||
|
||||
httpsReq, err := http.NewRequestWithContext(ctx, http.MethodGet, httpsURL, nil)
|
||||
if err != nil {
|
||||
return agent.ResultJSON(headersResult{
|
||||
ErrorDetail: fmt.Sprintf("cannot create request for %s: %s", httpsURL, err),
|
||||
}), nil
|
||||
}
|
||||
|
||||
resp, err := followClient.Do(httpsReq)
|
||||
if err != nil {
|
||||
return agent.ResultJSON(headersResult{
|
||||
ErrorDetail: fmt.Sprintf("cannot fetch %s: %s", httpsURL, err),
|
||||
}), nil
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
result := headersFromResponse(resp)
|
||||
|
||||
@@ -81,6 +81,7 @@ func CheckBreachesTool() agent.Tool {
|
||||
ErrorDetail: fmt.Sprintf("cannot fetch breaches: %s", err),
|
||||
}), nil
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
|
||||
@@ -26,6 +26,7 @@ func resolverAddr() string {
|
||||
if addr := os.Getenv("DNS_RESOLVER_ADDR"); addr != "" {
|
||||
return addr
|
||||
}
|
||||
|
||||
return "8.8.8.8:53"
|
||||
}
|
||||
|
||||
|
||||
@@ -65,6 +65,7 @@ func CheckSPFTool() agent.Tool {
|
||||
}
|
||||
|
||||
client := dns.NewClient()
|
||||
|
||||
answers, err := queryDNS(
|
||||
ctx,
|
||||
client,
|
||||
@@ -83,6 +84,7 @@ func CheckSPFTool() agent.Tool {
|
||||
}
|
||||
|
||||
var spfRecords []string
|
||||
|
||||
for _, answer := range answers {
|
||||
txt, ok := answer.(*dns.TXT)
|
||||
if !ok {
|
||||
@@ -106,6 +108,7 @@ func CheckSPFTool() agent.Tool {
|
||||
|
||||
if len(spfRecords) == 1 {
|
||||
record := spfRecords[0]
|
||||
|
||||
return agent.ResultJSON(spfResult{
|
||||
Found: true,
|
||||
RawRecord: record,
|
||||
|
||||
@@ -89,16 +89,19 @@ func CheckSSLCertificateTool() agent.Tool {
|
||||
},
|
||||
}
|
||||
netConn, err := dialer.DialContext(ctx, "tcp", p.Domain+":443")
|
||||
|
||||
var conn *tls.Conn
|
||||
if netConn != nil {
|
||||
conn = netConn.(*tls.Conn)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return agent.ResultJSON(sslResult{
|
||||
Valid: false,
|
||||
ErrorDetail: err.Error(),
|
||||
}), nil
|
||||
}
|
||||
|
||||
defer func() { _ = conn.Close() }()
|
||||
|
||||
state := conn.ConnectionState()
|
||||
@@ -124,6 +127,7 @@ func CheckSSLCertificateTool() agent.Tool {
|
||||
for _, ic := range state.PeerCertificates[1:] {
|
||||
opts.Intermediates.AddCert(ic)
|
||||
}
|
||||
|
||||
if _, err := cert.Verify(opts); err != nil {
|
||||
valid = false
|
||||
}
|
||||
|
||||
@@ -67,6 +67,7 @@ func CheckWhoisTool() agent.Tool {
|
||||
if whoisServer == "" {
|
||||
whoisServer = parseWhoisField(referral, "whois")
|
||||
}
|
||||
|
||||
if whoisServer == "" {
|
||||
// Try common TLD WHOIS servers as fallback.
|
||||
parts := strings.Split(p.Domain, ".")
|
||||
@@ -84,6 +85,7 @@ func CheckWhoisTool() agent.Tool {
|
||||
if whoisHost == "" {
|
||||
whoisHost = whoisServer
|
||||
}
|
||||
|
||||
if err := netcheck.ValidatePublicDomain(whoisHost); err != nil {
|
||||
return agent.ResultJSON(whoisResult{
|
||||
ErrorDetail: fmt.Sprintf("WHOIS referral server not allowed: %s", err),
|
||||
@@ -114,6 +116,7 @@ func CheckWhoisTool() agent.Tool {
|
||||
years := int(age.Hours() / 24 / 365)
|
||||
months := int(age.Hours()/24/30) % 12
|
||||
result.DomainAge = fmt.Sprintf("%d years, %d months", years, months)
|
||||
|
||||
break
|
||||
}
|
||||
}
|
||||
@@ -126,10 +129,12 @@ func CheckWhoisTool() agent.Tool {
|
||||
|
||||
func queryWhois(ctx context.Context, server, domain string) (string, error) {
|
||||
dialer := net.Dialer{Timeout: 10 * time.Second}
|
||||
|
||||
conn, err := dialer.DialContext(ctx, "tcp", server)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cannot connect to %s: %w", server, err)
|
||||
}
|
||||
|
||||
defer func() { _ = conn.Close() }()
|
||||
|
||||
_ = conn.SetDeadline(time.Now().Add(10 * time.Second))
|
||||
@@ -140,11 +145,13 @@ func queryWhois(ctx context.Context, server, domain string) (string, error) {
|
||||
}
|
||||
|
||||
var sb strings.Builder
|
||||
|
||||
scanner := bufio.NewScanner(conn)
|
||||
for scanner.Scan() {
|
||||
sb.WriteString(scanner.Text())
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
return "", fmt.Errorf("cannot read from %s: %w", server, err)
|
||||
}
|
||||
@@ -154,19 +161,23 @@ func queryWhois(ctx context.Context, server, domain string) (string, error) {
|
||||
|
||||
func parseWhoisField(raw, field string) string {
|
||||
field = strings.ToLower(field)
|
||||
|
||||
for line := range strings.SplitSeq(raw, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" || strings.HasPrefix(line, "%") || strings.HasPrefix(line, "#") {
|
||||
continue
|
||||
}
|
||||
|
||||
k, v, ok := strings.Cut(line, ":")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
if strings.ToLower(strings.TrimSpace(k)) == field {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
@@ -199,16 +210,20 @@ var (
|
||||
|
||||
func parseWhoisResponse(raw string) whoisResult {
|
||||
var result whoisResult
|
||||
|
||||
for line := range strings.SplitSeq(raw, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" || strings.HasPrefix(line, "%") || strings.HasPrefix(line, "#") {
|
||||
continue
|
||||
}
|
||||
|
||||
k, v, ok := strings.Cut(line, ":")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
key := strings.ToLower(strings.TrimSpace(k))
|
||||
|
||||
val := strings.TrimSpace(v)
|
||||
if val == "" {
|
||||
continue
|
||||
|
||||
@@ -33,8 +33,10 @@ func (m *typedMockProvider) ChatCompletion(_ context.Context, _ *llm.ChatComplet
|
||||
if m.calls >= len(m.responses) {
|
||||
return nil, errors.New("no more mock responses")
|
||||
}
|
||||
|
||||
resp := m.responses[m.calls]
|
||||
m.calls++
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -64,8 +64,10 @@ func (m *mockProvider) ChatCompletion(_ context.Context, _ *llm.ChatCompletionRe
|
||||
if m.calls >= len(m.responses) {
|
||||
return nil, errors.New("no more mock responses")
|
||||
}
|
||||
|
||||
resp := m.responses[m.calls]
|
||||
m.calls++
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
@@ -114,6 +116,7 @@ func (r *simpleRegistry) Agent(name string) (*agent.Agent, error) {
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("agent %q not found", name)
|
||||
}
|
||||
|
||||
return a, nil
|
||||
}
|
||||
|
||||
@@ -206,6 +209,7 @@ func TestAgentRunSupervisor_StopAndResume(t *testing.T) {
|
||||
func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
|
||||
close(toolReady)
|
||||
<-toolRelease
|
||||
|
||||
return agent.ToolResult{Content: "work done"}, nil
|
||||
},
|
||||
)
|
||||
@@ -332,6 +336,7 @@ func TestAgentRunSupervisor_StopAndResume(t *testing.T) {
|
||||
WHERE id = $1`,
|
||||
run.ID.String(),
|
||||
)
|
||||
|
||||
return err
|
||||
},
|
||||
)
|
||||
@@ -524,11 +529,13 @@ func makeBattleTools(progressFile string) []agent.Tool {
|
||||
// Record completion — written AFTER the sleep so the parent's
|
||||
// step count reflects truly-finished work.
|
||||
mu.Lock()
|
||||
|
||||
f, err := os.OpenFile(progressFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
|
||||
if err != nil {
|
||||
mu.Unlock()
|
||||
return agent.ToolResult{}, err
|
||||
}
|
||||
|
||||
_, _ = fmt.Fprintln(f, input.Task)
|
||||
_ = f.Close()
|
||||
mu.Unlock()
|
||||
@@ -584,12 +591,15 @@ func TestAgentRunSupervisor_SIGTERM(t *testing.T) {
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
|
||||
n := 0
|
||||
|
||||
for _, b := range data {
|
||||
if b == '\n' {
|
||||
n++
|
||||
}
|
||||
}
|
||||
|
||||
return n
|
||||
}
|
||||
|
||||
@@ -599,6 +609,7 @@ func TestAgentRunSupervisor_SIGTERM(t *testing.T) {
|
||||
"-test.run=^TestAgentRunSupervisor_SIGTERM$",
|
||||
"-test.v",
|
||||
)
|
||||
|
||||
cmd.Env = append(os.Environ(),
|
||||
"TEST_SIGTERM_SUBPROCESS=1",
|
||||
"TEST_SIGTERM_PROGRESS_FILE="+progressFile,
|
||||
@@ -607,27 +618,33 @@ func TestAgentRunSupervisor_SIGTERM(t *testing.T) {
|
||||
cmd.Stdout = os.Stdout
|
||||
cmd.Stderr = os.Stderr
|
||||
require.NoError(t, cmd.Start())
|
||||
|
||||
return cmd
|
||||
}
|
||||
|
||||
killAndWait := func(cmd *exec.Cmd) {
|
||||
require.NoError(t, cmd.Process.Signal(syscall.SIGTERM))
|
||||
|
||||
err := cmd.Wait()
|
||||
if err == nil {
|
||||
return
|
||||
}
|
||||
|
||||
exitErr, ok := errors.AsType[*exec.ExitError](err)
|
||||
if !ok {
|
||||
t.Fatalf("subprocess error: %v", err)
|
||||
}
|
||||
|
||||
ws, ok := exitErr.Sys().(syscall.WaitStatus)
|
||||
if !ok {
|
||||
t.Fatalf("subprocess exited with unexpected wait status: %v", exitErr)
|
||||
}
|
||||
|
||||
if ws.Signaled() && ws.Signal() == syscall.SIGTERM {
|
||||
t.Logf("subprocess terminated by SIGTERM")
|
||||
return
|
||||
}
|
||||
|
||||
t.Fatalf("subprocess exited unexpectedly (signaled=%v signal=%v exit=%d): %v",
|
||||
ws.Signaled(), ws.Signal(), ws.ExitStatus(), exitErr)
|
||||
}
|
||||
@@ -646,6 +663,7 @@ func TestAgentRunSupervisor_SIGTERM(t *testing.T) {
|
||||
WHERE id = $1`,
|
||||
run.ID.String(),
|
||||
)
|
||||
|
||||
return err
|
||||
},
|
||||
)
|
||||
@@ -674,6 +692,7 @@ func TestAgentRunSupervisor_SIGTERM(t *testing.T) {
|
||||
" checkpoint: %d messages, %d turns, usage=%+v",
|
||||
len(cp.Messages), cp.Turns, cp.Usage,
|
||||
)
|
||||
|
||||
return cp
|
||||
}
|
||||
|
||||
@@ -682,7 +701,9 @@ func TestAgentRunSupervisor_SIGTERM(t *testing.T) {
|
||||
// Steps so far: turn0=1(scan) + turn1=3(fetch×3) = 4
|
||||
// ============================================================
|
||||
t.Log("=== Phase 1: SIGTERM after scan + parallel fetch (4 steps) ===")
|
||||
|
||||
cmd1 := startSubprocess(0)
|
||||
|
||||
waitForSteps(4)
|
||||
killAndWait(cmd1)
|
||||
|
||||
@@ -700,7 +721,9 @@ func TestAgentRunSupervisor_SIGTERM(t *testing.T) {
|
||||
// New steps: turn2=1(analyze) + turn3=3(check×3) = 4
|
||||
// ============================================================
|
||||
t.Log("=== Phase 2: SIGTERM after analyze + parallel checks (4 more steps) ===")
|
||||
|
||||
cmd2 := startSubprocess(cp1.Turns)
|
||||
|
||||
waitForSteps(steps1 + 4)
|
||||
killAndWait(cmd2)
|
||||
|
||||
@@ -720,7 +743,9 @@ func TestAgentRunSupervisor_SIGTERM(t *testing.T) {
|
||||
// New steps: turn4=1(deep) + turn5=2(generate×2) = 3
|
||||
// ============================================================
|
||||
t.Log("=== Phase 3: SIGTERM during long-running deep analysis (3 more steps) ===")
|
||||
|
||||
cmd3 := startSubprocess(cp2.Turns)
|
||||
|
||||
waitForSteps(steps2 + 3)
|
||||
killAndWait(cmd3)
|
||||
|
||||
@@ -811,14 +836,17 @@ func runSIGTERMSubprocess() {
|
||||
if addr == "" {
|
||||
addr = "localhost:5432"
|
||||
}
|
||||
|
||||
user := os.Getenv("PROBO_TEST_PG_USER")
|
||||
if user == "" {
|
||||
user = "probod"
|
||||
}
|
||||
|
||||
password := os.Getenv("PROBO_TEST_PG_PASSWORD")
|
||||
if password == "" {
|
||||
password = "probod"
|
||||
}
|
||||
|
||||
database := os.Getenv("PROBO_TEST_PG_DATABASE")
|
||||
if database == "" {
|
||||
database = "probod_test"
|
||||
|
||||
@@ -115,6 +115,7 @@ func EnsureAgentRunsTable(t *testing.T, client *pg.Client) {
|
||||
).Scan(&exists); err != nil {
|
||||
return fmt.Errorf("cannot check agent_runs existence: %w", err)
|
||||
}
|
||||
|
||||
if exists {
|
||||
return nil
|
||||
}
|
||||
@@ -127,6 +128,7 @@ func EnsureAgentRunsTable(t *testing.T, client *pg.Client) {
|
||||
if _, err := conn.Exec(ctx, string(ddl)); err != nil {
|
||||
return fmt.Errorf("cannot apply agent_runs migration: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
})
|
||||
@@ -185,6 +187,7 @@ func InsertPendingRun(
|
||||
); err != nil {
|
||||
return fmt.Errorf("cannot insert placeholder organization: %w", err)
|
||||
}
|
||||
|
||||
return run.Insert(ctx, tx, coredata.NewScope(tenantID))
|
||||
},
|
||||
)
|
||||
@@ -214,6 +217,7 @@ func LoadAgentRun(t *testing.T, client *pg.Client, id gid.GID) coredata.AgentRun
|
||||
t.Helper()
|
||||
|
||||
var run coredata.AgentRun
|
||||
|
||||
err := client.WithConn(
|
||||
context.Background(),
|
||||
func(ctx context.Context, conn pg.Querier) error {
|
||||
@@ -231,11 +235,13 @@ func LoadAgentRun(t *testing.T, client *pg.Client, id gid.GID) coredata.AgentRun
|
||||
// require.Eventually callbacks (which recover panics).
|
||||
func TryLoadAgentRun(client *pg.Client, id gid.GID) (coredata.AgentRun, error) {
|
||||
var run coredata.AgentRun
|
||||
|
||||
err := client.WithConn(
|
||||
context.Background(),
|
||||
func(ctx context.Context, conn pg.Querier) error {
|
||||
return run.LoadByID(ctx, conn, coredata.NewNoScope(), id)
|
||||
},
|
||||
)
|
||||
|
||||
return run, err
|
||||
}
|
||||
|
||||
@@ -64,6 +64,7 @@ func MustParse(rawURL string) *BaseURL {
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
return b
|
||||
}
|
||||
|
||||
@@ -72,6 +73,7 @@ func (b *BaseURL) String() string {
|
||||
if b == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
return b.raw
|
||||
}
|
||||
|
||||
@@ -80,6 +82,7 @@ func (b *BaseURL) Scheme() string {
|
||||
if b == nil || b.parsed == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
return b.parsed.Scheme
|
||||
}
|
||||
|
||||
@@ -88,6 +91,7 @@ func (b *BaseURL) Host() string {
|
||||
if b == nil || b.parsed == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
return b.parsed.Host
|
||||
}
|
||||
|
||||
@@ -96,6 +100,7 @@ func (b *BaseURL) Hostname() string {
|
||||
if b == nil || b.parsed == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
return b.parsed.Hostname()
|
||||
}
|
||||
|
||||
@@ -104,6 +109,7 @@ func (b *BaseURL) Port() string {
|
||||
if b == nil || b.parsed == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
return b.parsed.Port()
|
||||
}
|
||||
|
||||
@@ -168,7 +174,9 @@ func (ub *URLBuilder) WithQuery(key, value string) *URLBuilder {
|
||||
if ub.err != nil {
|
||||
return ub
|
||||
}
|
||||
|
||||
ub.query.Add(key, value)
|
||||
|
||||
return ub
|
||||
}
|
||||
|
||||
@@ -177,11 +185,13 @@ func (ub *URLBuilder) WithQueryValues(values url.Values) *URLBuilder {
|
||||
if ub.err != nil {
|
||||
return ub
|
||||
}
|
||||
|
||||
for key, vals := range values {
|
||||
for _, val := range vals {
|
||||
ub.query.Add(key, val)
|
||||
}
|
||||
}
|
||||
|
||||
return ub
|
||||
}
|
||||
|
||||
@@ -208,6 +218,7 @@ func (ub *URLBuilder) MustString() string {
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
@@ -224,6 +235,7 @@ func (b *BaseURL) UnmarshalJSON(data []byte) error {
|
||||
}
|
||||
|
||||
*b = *parsed
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -232,6 +244,7 @@ func (b *BaseURL) MarshalJSON() ([]byte, error) {
|
||||
if b == nil {
|
||||
return json.Marshal("")
|
||||
}
|
||||
|
||||
return json.Marshal(b.raw)
|
||||
}
|
||||
|
||||
@@ -243,6 +256,7 @@ func (b *BaseURL) UnmarshalText(text []byte) error {
|
||||
}
|
||||
|
||||
*b = *parsed
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -251,6 +265,7 @@ func (b *BaseURL) MarshalText() ([]byte, error) {
|
||||
if b == nil {
|
||||
return []byte(""), nil
|
||||
}
|
||||
|
||||
return []byte(b.raw), nil
|
||||
}
|
||||
|
||||
|
||||
@@ -75,6 +75,7 @@ func TestParse(t *testing.T) {
|
||||
t.Errorf("Parse() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
|
||||
if !tt.wantErr && got == nil {
|
||||
t.Error("Parse() returned nil without error")
|
||||
}
|
||||
@@ -133,6 +134,7 @@ func TestBaseURL_WithPath(t *testing.T) {
|
||||
t.Errorf("WithPath().String() error = %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if got != tt.want {
|
||||
t.Errorf("WithPath().String() = %v, want %v", got, tt.want)
|
||||
}
|
||||
@@ -147,7 +149,6 @@ func TestBaseURL_WithQuery(t *testing.T) {
|
||||
WithQuery("q", "test").
|
||||
WithQuery("limit", "10").
|
||||
String()
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("WithPath().WithQuery().String() error = %v", err)
|
||||
}
|
||||
|
||||
@@ -36,6 +36,7 @@ func NewBuilder(getEnv EnvGetter) *Builder {
|
||||
if getEnv == nil {
|
||||
getEnv = os.Getenv
|
||||
}
|
||||
|
||||
return &Builder{getEnv: getEnv}
|
||||
}
|
||||
|
||||
@@ -372,6 +373,7 @@ func (b *Builder) Build() (*probodconfig.FullConfig, error) {
|
||||
if clientID == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
cfg.Probod.Connectors = append(cfg.Probod.Connectors, probodconfig.ConnectorConfig{
|
||||
Provider: provider,
|
||||
Protocol: "oauth2",
|
||||
@@ -480,6 +482,7 @@ func (b *Builder) getSAMLCredentials() (cert, key string, err error) {
|
||||
if cert == "" {
|
||||
cert = b.getEnv("SAML_CERTIFICATE")
|
||||
}
|
||||
|
||||
if key == "" {
|
||||
key = b.getEnv("SAML_PRIVATE_KEY")
|
||||
}
|
||||
@@ -498,6 +501,7 @@ func (b *Builder) getOAuth2SigningKey() string {
|
||||
if b.oauth2SigningKey != "" {
|
||||
return b.oauth2SigningKey
|
||||
}
|
||||
|
||||
return b.getEnv("OAUTH2_SERVER_SIGNING_KEY")
|
||||
}
|
||||
|
||||
@@ -516,6 +520,7 @@ func (b *Builder) getEnvOrDefault(key, defaultValue string) string {
|
||||
if value := b.getEnv(key); value != "" {
|
||||
return value
|
||||
}
|
||||
|
||||
return defaultValue
|
||||
}
|
||||
|
||||
@@ -525,6 +530,7 @@ func (b *Builder) getEnvIntOrDefault(key string, defaultValue int) int {
|
||||
return int(intValue)
|
||||
}
|
||||
}
|
||||
|
||||
return defaultValue
|
||||
}
|
||||
|
||||
@@ -534,6 +540,7 @@ func (b *Builder) getEnvFloatOrDefault(key string, defaultValue float64) float64
|
||||
return floatValue
|
||||
}
|
||||
}
|
||||
|
||||
return defaultValue
|
||||
}
|
||||
|
||||
@@ -543,6 +550,7 @@ func (b *Builder) getEnvFloatPtr(key string) *float64 {
|
||||
return &floatValue
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -553,6 +561,7 @@ func (b *Builder) getEnvIntPtr(key string) *int {
|
||||
return &v
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -562,6 +571,7 @@ func (b *Builder) getEnvBoolOrDefault(key string, defaultValue bool) bool {
|
||||
return boolValue
|
||||
}
|
||||
}
|
||||
|
||||
return defaultValue
|
||||
}
|
||||
|
||||
@@ -572,12 +582,15 @@ func (b *Builder) parseOriginsList(s string) []string {
|
||||
}
|
||||
|
||||
var result []string
|
||||
|
||||
for part := range strings.SplitSeq(s, ",") {
|
||||
part = strings.TrimSpace(part)
|
||||
|
||||
part = strings.Trim(part, "\"")
|
||||
if part != "" {
|
||||
result = append(result, part)
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
@@ -114,6 +114,7 @@ func TestBuilder_Build_MissingRequiredEnvVars(t *testing.T) {
|
||||
_, err := b.Build()
|
||||
|
||||
require.Error(t, err)
|
||||
|
||||
for _, missing := range tt.wantMissing {
|
||||
assert.Contains(t, err.Error(), missing)
|
||||
}
|
||||
@@ -480,6 +481,7 @@ func TestBuilder_Build_AccessReviewConnectors(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Len(t, cfg.Probod.Connectors, len(providers))
|
||||
|
||||
byProvider := make(map[string]probodconfig.ConnectorConfig, len(cfg.Probod.Connectors))
|
||||
for _, c := range cfg.Probod.Connectors {
|
||||
byProvider[c.Provider] = c
|
||||
@@ -539,6 +541,7 @@ func TestBuilder_Build_SlackConnector(t *testing.T) {
|
||||
rawConfig := connector.RawConfig.(probodconfig.ConnectorConfigOAuth2)
|
||||
assert.Equal(t, "slack-client-id", rawConfig.ClientID)
|
||||
assert.Equal(t, "slack-client-secret", rawConfig.ClientSecret)
|
||||
|
||||
rawSettings := connector.RawSettings.(map[string]any)
|
||||
assert.Equal(t, "slack-signing-secret", rawSettings["signing-secret"])
|
||||
}
|
||||
|
||||
@@ -46,6 +46,7 @@ func TestWriteConfig(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
var loaded probodconfig.FullConfig
|
||||
|
||||
err = yaml.Unmarshal(data, &loaded)
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -143,6 +144,7 @@ func TestWriteConfig_CompleteConfig(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
var loaded probodconfig.FullConfig
|
||||
|
||||
err = yaml.Unmarshal(data, &loaded)
|
||||
require.NoError(t, err)
|
||||
|
||||
|
||||
@@ -66,6 +66,7 @@ func NewACMEService(
|
||||
) (*ACMEService, error) {
|
||||
if accountKey == nil {
|
||||
var err error
|
||||
|
||||
accountKey, err = keys.Generate(keyType)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot generate account key: %w", err)
|
||||
@@ -130,6 +131,7 @@ func (s *ACMEService) GetHTTPChallenge(ctx context.Context, domain string) (*HTT
|
||||
}
|
||||
|
||||
var challenge *acme.Challenge
|
||||
|
||||
for _, auth := range order.AuthzURLs {
|
||||
authz, err := s.client.GetAuthorization(ctx, auth)
|
||||
if err != nil {
|
||||
@@ -170,7 +172,6 @@ func (s *ACMEService) CompleteHTTPChallenge(
|
||||
ctx context.Context,
|
||||
challenge0 *HTTPChallenge,
|
||||
) (*Certificate, error) {
|
||||
|
||||
challenge1 := &acme.Challenge{
|
||||
URI: challenge0.URL,
|
||||
Token: challenge0.Token,
|
||||
@@ -206,6 +207,7 @@ func (s *ACMEService) CompleteHTTPChallenge(
|
||||
}
|
||||
|
||||
certPEM := pem.EncodeCertificate(der[0])
|
||||
|
||||
keyPEM, err := pem.EncodePrivateKey(certKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot encode key: %w", err)
|
||||
@@ -215,6 +217,7 @@ func (s *ACMEService) CompleteHTTPChallenge(
|
||||
if len(der) > 1 {
|
||||
chainDER = der[1:]
|
||||
}
|
||||
|
||||
chainPEM := pem.EncodeCertificateChain(chainDER)
|
||||
|
||||
return &Certificate{
|
||||
|
||||
@@ -62,6 +62,7 @@ func (h *ACMEChallengeHandler) Handle(next http.Handler) http.Handler {
|
||||
)
|
||||
|
||||
http.NotFound(w, r)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
@@ -47,6 +47,7 @@ func NewCacheStore(
|
||||
|
||||
func (w *CacheStore) WarmCache(ctx context.Context) error {
|
||||
w.logger.InfoCtx(ctx, "warming certificate cache")
|
||||
|
||||
startTime := time.Now()
|
||||
|
||||
err := w.pg.WithConn(
|
||||
@@ -65,6 +66,7 @@ func (w *CacheStore) WarmCache(ctx context.Context) error {
|
||||
w.logger.InfoCtx(ctx, "found active certificates to cache", log.Int("count", len(domains)))
|
||||
|
||||
successCount := 0
|
||||
|
||||
for _, domain := range domains {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
@@ -80,10 +82,10 @@ func (w *CacheStore) WarmCache(ctx context.Context) error {
|
||||
}
|
||||
|
||||
w.logger.InfoCtx(ctx, "successfully warmed cache", log.Int("success_count", successCount), log.Int("total_count", len(domains)))
|
||||
|
||||
return nil
|
||||
},
|
||||
)
|
||||
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot warm certificate cache: %w", err)
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user