Add wsl linter and fix

Signed-off-by: Émile Ré <emile@probo.com>
This commit is contained in:
Émile Ré
2026-05-19 14:51:08 +04:00
parent eedfdcecc8
commit 9156d6a16a
882 changed files with 6068 additions and 574 deletions

View File

@@ -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
},
)

View File

@@ -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 {

View File

@@ -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)

View File

@@ -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 {

View File

@@ -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 {

View File

@@ -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()
}()

View File

@@ -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

View File

@@ -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
}

View File

@@ -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()
}()

View File

@@ -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

View File

@@ -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)
}

View File

@@ -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()
}()

View File

@@ -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()
}()

View File

@@ -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 {

View File

@@ -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 {

View File

@@ -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

View File

@@ -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()
}()

View File

@@ -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"
}

View File

@@ -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)
}

View File

@@ -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 {

View File

@@ -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 {

View File

@@ -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
}

View File

@@ -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)
})

View File

@@ -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 {

View File

@@ -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()
}()

View File

@@ -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()
}()

View File

@@ -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()
}()

View File

@@ -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()
}()

View File

@@ -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
}

View File

@@ -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 {

View File

@@ -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()
}()

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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)

View File

@@ -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()
}()

View File

@@ -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,

View File

@@ -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}
}

View File

@@ -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 {

View File

@@ -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

View File

@@ -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
},
)

View File

@@ -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)

View File

@@ -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)
},
)

View File

@@ -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
}

View File

@@ -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",

View File

@@ -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,

View File

@@ -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())

View File

@@ -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
}

View File

@@ -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
},
}

View File

@@ -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)

View File

@@ -49,5 +49,6 @@ func TryRunContextFrom[C any](ctx context.Context) (C, bool) {
}
typed, ok := val.(C)
return typed, ok
}

View File

@@ -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
}

View File

@@ -32,6 +32,7 @@ func NewSystemPromptLeakGuardrail(fingerprints []string) *SystemPromptLeakGuardr
if f == "" {
continue
}
lowered = append(lowered, strings.ToLower(f))
}

View File

@@ -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)
}

View File

@@ -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
}),
)

View File

@@ -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)

View File

@@ -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"},

View File

@@ -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

View File

@@ -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
}
}

View File

@@ -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)

View File

@@ -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]
}

View File

@@ -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
}

View File

@@ -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

View File

@@ -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"`
}

View File

@@ -61,6 +61,7 @@ func (s *memorySession) Save(_ context.Context, sessionID string, messages []llm
}
s.sessions[sessionID] = cp
return nil
}

View File

@@ -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()

View File

@@ -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)

View File

@@ -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(

View File

@@ -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
}
}

View File

@@ -123,6 +123,7 @@ func (b *Browser) checkAlive() *agent.ToolResult {
IsError: true,
}
}
return nil
}

View File

@@ -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")
}

View File

@@ -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)

View File

@@ -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
}

View File

@@ -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)
},
}

View File

@@ -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++
}

View File

@@ -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)

View File

@@ -91,6 +91,7 @@ func CheckGovernmentDBTool(apiKey string) agent.Tool {
if err != nil {
continue
}
for _, e := range entries {
*s.target = append(
*s.target,

View File

@@ -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
}

View File

@@ -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 {

View File

@@ -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")

View File

@@ -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 {

View File

@@ -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,

View File

@@ -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
}

View File

@@ -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),

View File

@@ -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)

View File

@@ -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)

View File

@@ -26,6 +26,7 @@ func resolverAddr() string {
if addr := os.Getenv("DNS_RESOLVER_ADDR"); addr != "" {
return addr
}
return "8.8.8.8:53"
}

View File

@@ -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,

View File

@@ -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
}

View File

@@ -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

View File

@@ -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
}

View File

@@ -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"

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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)
}

View File

@@ -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
}

View File

@@ -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"])
}

View File

@@ -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)

View File

@@ -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{

View File

@@ -62,6 +62,7 @@ func (h *ACMEChallengeHandler) Handle(next http.Handler) http.Handler {
)
http.NotFound(w, r)
return
}

View File

@@ -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