From 49d9a96355f34d399ef2363fe0d896bf07c0ad93 Mon Sep 17 00:00:00 2001 From: Bryan Frimin Date: Mon, 8 Jun 2026 07:32:36 +0200 Subject: [PATCH] Style Signed-off-by: Bryan Frimin --- e2e/console/agent_run_test.go | 480 ++++++++++++++++++++++++++++++++++ pkg/agent/agent_tool_test.go | 3 + pkg/agentrun/handler.go | 2 + pkg/agentrun/helpers_test.go | 2 + pkg/agentrun/worker_test.go | 45 +++- 5 files changed, 529 insertions(+), 3 deletions(-) create mode 100644 e2e/console/agent_run_test.go diff --git a/e2e/console/agent_run_test.go b/e2e/console/agent_run_test.go new file mode 100644 index 000000000..399a7eac9 --- /dev/null +++ b/e2e/console/agent_run_test.go @@ -0,0 +1,480 @@ +// Copyright (c) 2026 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package console_test + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.probo.inc/probo/e2e/internal/testutil" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/gid" +) + +// agentRunSeed describes the agent run row inserted directly into the test +// database. Agent runs have no creation mutation on the console API (they are +// produced by the agent run worker), so e2e coverage seeds them straight into +// Postgres, mirroring the common-third-party catalog seeding helper. +type agentRunSeed struct { + agentName string + status coredata.AgentRunStatus + errorMessage *string + startedAt *time.Time + createdAt time.Time +} + +func seedAgentRun(t *testing.T, organizationID gid.GID, seed agentRunSeed) gid.GID { + t.Helper() + + ctx := context.Background() + conn := dialTestPg(t, ctx) + t.Cleanup(func() { _ = conn.Close(ctx) }) + + if seed.agentName == "" { + seed.agentName = "test-agent" + } + + if seed.status == "" { + seed.status = coredata.AgentRunStatusPending + } + + if seed.createdAt.IsZero() { + seed.createdAt = time.Now().UTC() + } + + id := gid.New(organizationID.TenantID(), coredata.AgentRunEntityType) + + _, err := conn.Exec(ctx, ` + INSERT INTO agent_runs ( + id, tenant_id, organization_id, start_agent_name, status, + input_messages, error_message, started_at, created_at, updated_at + ) VALUES ( + $1, $2, $3, $4, $5, $6, $7, $8, $9, $9 + ) + `, + id, + organizationID.TenantID(), + organizationID, + seed.agentName, + seed.status, + []byte(`[]`), + seed.errorMessage, + seed.startedAt, + seed.createdAt, + ) + require.NoError(t, err, "cannot seed agent run") + + t.Cleanup(func() { + cleanupCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + cleanupConn := dialTestPg(t, cleanupCtx) + + defer func() { _ = cleanupConn.Close(cleanupCtx) }() + + _, err := cleanupConn.Exec(cleanupCtx, `DELETE FROM agent_runs WHERE id = $1`, id) + assert.NoError(t, err, "cleanup: cannot delete seeded agent run %s", id) + }) + + return id +} + +const agentRunListQuery = ` + query($orgId: ID!, $orderBy: AgentRunOrder) { + node(id: $orgId) { + ... on Organization { + agentRuns(first: 50, orderBy: $orderBy) { + totalCount + edges { + cursor + node { + id + agentName + status + errorMessage + startedAt + createdAt + updatedAt + } + } + pageInfo { + hasNextPage + hasPreviousPage + startCursor + endCursor + } + } + } + } + } +` + +type agentRunNode struct { + ID string `json:"id"` + AgentName string `json:"agentName"` + Status string `json:"status"` + ErrorMessage *string `json:"errorMessage"` + StartedAt *string `json:"startedAt"` + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` +} + +type agentRunConnectionResult struct { + Node struct { + AgentRuns struct { + TotalCount int `json:"totalCount"` + Edges []struct { + Cursor string `json:"cursor"` + Node agentRunNode `json:"node"` + } `json:"edges"` + PageInfo testutil.PageInfo `json:"pageInfo"` + } `json:"agentRuns"` + } `json:"node"` +} + +func TestAgentRun_List(t *testing.T) { + t.Parallel() + owner := testutil.NewClient(t, testutil.RoleOwner) + + errMsg := "boom" + startedAt := time.Now().UTC().Add(-time.Minute) + + completedID := seedAgentRun(t, owner.GetOrganizationID(), agentRunSeed{ + agentName: "compliance-agent", + status: coredata.AgentRunStatusCompleted, + startedAt: &startedAt, + }) + failedID := seedAgentRun(t, owner.GetOrganizationID(), agentRunSeed{ + agentName: "vetting-agent", + status: coredata.AgentRunStatusFailed, + errorMessage: &errMsg, + startedAt: &startedAt, + }) + + var result agentRunConnectionResult + + err := owner.Execute(agentRunListQuery, map[string]any{ + "orgId": owner.GetOrganizationID().String(), + }, &result) + require.NoError(t, err) + + assert.Equal(t, 2, result.Node.AgentRuns.TotalCount) + require.Len(t, result.Node.AgentRuns.Edges, 2) + + byID := make(map[string]agentRunNode, 2) + + for _, edge := range result.Node.AgentRuns.Edges { + assert.NotEmpty(t, edge.Cursor, "edge cursor should be set") + byID[edge.Node.ID] = edge.Node + } + + completed, ok := byID[completedID.String()] + require.True(t, ok, "completed run not returned in list") + assert.Equal(t, "compliance-agent", completed.AgentName) + assert.Equal(t, "COMPLETED", completed.Status) + assert.Nil(t, completed.ErrorMessage) + assert.NotNil(t, completed.StartedAt) + assert.NotEmpty(t, completed.CreatedAt) + assert.NotEmpty(t, completed.UpdatedAt) + + failed, ok := byID[failedID.String()] + require.True(t, ok, "failed run not returned in list") + assert.Equal(t, "vetting-agent", failed.AgentName) + assert.Equal(t, "FAILED", failed.Status) + require.NotNil(t, failed.ErrorMessage) + assert.Equal(t, "boom", *failed.ErrorMessage) +} + +func TestAgentRun_ListEmpty(t *testing.T) { + t.Parallel() + owner := testutil.NewClient(t, testutil.RoleOwner) + + var result agentRunConnectionResult + + err := owner.Execute(agentRunListQuery, map[string]any{ + "orgId": owner.GetOrganizationID().String(), + }, &result) + require.NoError(t, err) + + assert.Equal(t, 0, result.Node.AgentRuns.TotalCount) + assert.Empty(t, result.Node.AgentRuns.Edges) + assert.False(t, result.Node.AgentRuns.PageInfo.HasNextPage) + assert.False(t, result.Node.AgentRuns.PageInfo.HasPreviousPage) +} + +func TestAgentRun_Ordering(t *testing.T) { + t.Parallel() + owner := testutil.NewClient(t, testutil.RoleOwner) + + base := time.Now().UTC().Add(-time.Hour) + oldestID := seedAgentRun(t, owner.GetOrganizationID(), agentRunSeed{createdAt: base}) + middleID := seedAgentRun(t, owner.GetOrganizationID(), agentRunSeed{createdAt: base.Add(time.Minute)}) + newestID := seedAgentRun(t, owner.GetOrganizationID(), agentRunSeed{createdAt: base.Add(2 * time.Minute)}) + + t.Run("ascending by createdAt", func(t *testing.T) { + t.Parallel() + + var result agentRunConnectionResult + + err := owner.Execute(agentRunListQuery, map[string]any{ + "orgId": owner.GetOrganizationID().String(), + "orderBy": map[string]any{"direction": "ASC", "field": "CREATED_AT"}, + }, &result) + require.NoError(t, err) + require.Len(t, result.Node.AgentRuns.Edges, 3) + + assert.Equal(t, oldestID.String(), result.Node.AgentRuns.Edges[0].Node.ID) + assert.Equal(t, middleID.String(), result.Node.AgentRuns.Edges[1].Node.ID) + assert.Equal(t, newestID.String(), result.Node.AgentRuns.Edges[2].Node.ID) + }) + + t.Run("descending by createdAt", func(t *testing.T) { + t.Parallel() + + var result agentRunConnectionResult + + err := owner.Execute(agentRunListQuery, map[string]any{ + "orgId": owner.GetOrganizationID().String(), + "orderBy": map[string]any{"direction": "DESC", "field": "CREATED_AT"}, + }, &result) + require.NoError(t, err) + require.Len(t, result.Node.AgentRuns.Edges, 3) + + assert.Equal(t, newestID.String(), result.Node.AgentRuns.Edges[0].Node.ID) + assert.Equal(t, middleID.String(), result.Node.AgentRuns.Edges[1].Node.ID) + assert.Equal(t, oldestID.String(), result.Node.AgentRuns.Edges[2].Node.ID) + }) +} + +func TestAgentRun_Pagination(t *testing.T) { + t.Parallel() + owner := testutil.NewClient(t, testutil.RoleOwner) + + base := time.Now().UTC().Add(-time.Hour) + for i := range 3 { + seedAgentRun(t, owner.GetOrganizationID(), agentRunSeed{ + createdAt: base.Add(time.Duration(i) * time.Minute), + }) + } + + const query = ` + query($orgId: ID!, $first: Int, $after: CursorKey) { + node(id: $orgId) { + ... on Organization { + agentRuns(first: $first, after: $after, orderBy: {direction: ASC, field: CREATED_AT}) { + totalCount + edges { + cursor + node { id } + } + pageInfo { + hasNextPage + hasPreviousPage + startCursor + endCursor + } + } + } + } + } + ` + + var firstPage agentRunConnectionResult + + err := owner.Execute(query, map[string]any{ + "orgId": owner.GetOrganizationID().String(), + "first": 2, + }, &firstPage) + require.NoError(t, err) + + assert.Equal(t, 3, firstPage.Node.AgentRuns.TotalCount) + testutil.AssertFirstPage(t, len(firstPage.Node.AgentRuns.Edges), firstPage.Node.AgentRuns.PageInfo, 2, true) + require.NotNil(t, firstPage.Node.AgentRuns.PageInfo.EndCursor) + + var secondPage agentRunConnectionResult + + err = owner.Execute(query, map[string]any{ + "orgId": owner.GetOrganizationID().String(), + "first": 2, + "after": *firstPage.Node.AgentRuns.PageInfo.EndCursor, + }, &secondPage) + require.NoError(t, err) + + testutil.AssertLastPage(t, len(secondPage.Node.AgentRuns.Edges), secondPage.Node.AgentRuns.PageInfo, 1, true) + + // The page boundary must not overlap. + firstIDs := map[string]struct{}{} + for _, edge := range firstPage.Node.AgentRuns.Edges { + firstIDs[edge.Node.ID] = struct{}{} + } + + for _, edge := range secondPage.Node.AgentRuns.Edges { + _, overlap := firstIDs[edge.Node.ID] + assert.False(t, overlap, "second page must not repeat a first-page run") + } +} + +func TestAgentRun_Get(t *testing.T) { + t.Parallel() + owner := testutil.NewClient(t, testutil.RoleOwner) + + runID := seedAgentRun(t, owner.GetOrganizationID(), agentRunSeed{ + agentName: "compliance-agent", + status: coredata.AgentRunStatusRunning, + }) + + const query = ` + query($id: ID!) { + node(id: $id) { + ... on AgentRun { + id + agentName + status + errorMessage + startedAt + createdAt + updatedAt + organization { id } + permission(action: "core:agent-run:get") + } + } + } + ` + + var result struct { + Node struct { + ID string `json:"id"` + AgentName string `json:"agentName"` + Status string `json:"status"` + ErrorMessage *string `json:"errorMessage"` + StartedAt *string `json:"startedAt"` + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` + Organization struct { + ID string `json:"id"` + } `json:"organization"` + Permission bool `json:"permission"` + } `json:"node"` + } + + err := owner.Execute(query, map[string]any{"id": runID.String()}, &result) + require.NoError(t, err) + + assert.Equal(t, runID.String(), result.Node.ID) + assert.Equal(t, "compliance-agent", result.Node.AgentName) + assert.Equal(t, "RUNNING", result.Node.Status) + assert.Nil(t, result.Node.ErrorMessage) + assert.Equal(t, owner.GetOrganizationID().String(), result.Node.Organization.ID) + assert.True(t, result.Node.Permission, "owner should have agent-run:get permission") +} + +func TestAgentRun_RBAC(t *testing.T) { + t.Parallel() + owner := testutil.NewClient(t, testutil.RoleOwner) + + runID := seedAgentRun(t, owner.GetOrganizationID(), agentRunSeed{ + agentName: "compliance-agent", + status: coredata.AgentRunStatusCompleted, + }) + + const getQuery = ` + query($id: ID!) { + node(id: $id) { + ... on AgentRun { + id + agentName + } + } + } + ` + + roles := []testutil.TestRole{testutil.RoleAdmin, testutil.RoleViewer} + for _, role := range roles { + t.Run(string(role)+" can list and get agent runs", func(t *testing.T) { + t.Parallel() + member := testutil.NewClientInOrg(t, role, owner) + + var listResult agentRunConnectionResult + + err := member.Execute(agentRunListQuery, map[string]any{ + "orgId": member.GetOrganizationID().String(), + }, &listResult) + require.NoError(t, err) + assert.Equal(t, 1, listResult.Node.AgentRuns.TotalCount) + + var getResult struct { + Node struct { + ID string `json:"id"` + AgentName string `json:"agentName"` + } `json:"node"` + } + + err = member.Execute(getQuery, map[string]any{"id": runID.String()}, &getResult) + require.NoError(t, err) + assert.Equal(t, runID.String(), getResult.Node.ID) + assert.Equal(t, "compliance-agent", getResult.Node.AgentName) + }) + } +} + +func TestAgentRun_TenantIsolation(t *testing.T) { + t.Parallel() + + org1Owner := testutil.NewClient(t, testutil.RoleOwner) + org2Owner := testutil.NewClient(t, testutil.RoleOwner) + + runID := seedAgentRun(t, org1Owner.GetOrganizationID(), agentRunSeed{ + agentName: "compliance-agent", + status: coredata.AgentRunStatusCompleted, + }) + + t.Run("other org cannot fetch the run by id", func(t *testing.T) { + t.Parallel() + + const query = ` + query($id: ID!) { + node(id: $id) { + ... on AgentRun { id } + } + } + ` + + var result struct { + Node *struct { + ID string `json:"id"` + } `json:"node"` + } + + err := org2Owner.Execute(query, map[string]any{"id": runID.String()}, &result) + testutil.AssertNodeNotAccessible(t, err, result.Node == nil, "AgentRun") + }) + + t.Run("other org list does not include the run", func(t *testing.T) { + t.Parallel() + + var result agentRunConnectionResult + + err := org2Owner.Execute(agentRunListQuery, map[string]any{ + "orgId": org2Owner.GetOrganizationID().String(), + }, &result) + require.NoError(t, err) + + assert.Equal(t, 0, result.Node.AgentRuns.TotalCount) + assert.Empty(t, result.Node.AgentRuns.Edges) + }) +} diff --git a/pkg/agent/agent_tool_test.go b/pkg/agent/agent_tool_test.go index 463608a57..c546fcd55 100644 --- a/pkg/agent/agent_tool_test.go +++ b/pkg/agent/agent_tool_test.go @@ -1210,6 +1210,7 @@ func TestAgentTool_Execute_LeafToolsRemainDetachedOnSuspend(t *testing.T) { if ctx.Err() != nil { leafCtxCanceled.Store(true) } + return agent.ToolResult{Content: "leaf completed"}, nil } }, @@ -1235,10 +1236,12 @@ func TestAgentTool_Execute_LeafToolsRemainDetachedOnSuspend(t *testing.T) { ) ctx, cancel := context.WithCancel(context.Background()) + type runResult struct { result *agent.Result err error } + runDone := make(chan runResult, 1) go func() { diff --git a/pkg/agentrun/handler.go b/pkg/agentrun/handler.go index 8bd79eacc..86b41f26f 100644 --- a/pkg/agentrun/handler.go +++ b/pkg/agentrun/handler.go @@ -100,6 +100,7 @@ func (h *handler) Claim(ctx context.Context) (coredata.AgentRun, error) { func (h *handler) Process(ctx context.Context, run coredata.AgentRun) error { runCtx, cancelRun := context.WithCancelCause(ctx) defer cancelRun(nil) + leaseGeneration := run.LeaseGeneration forwarderDone := make(chan struct{}) @@ -368,6 +369,7 @@ func (h *handler) executeRun( } h.logger.ErrorCtx(commitCtx, "cannot commit agent run status", log.Error(err)) + return fmt.Errorf("cannot commit agent run status: %w", err) } diff --git a/pkg/agentrun/helpers_test.go b/pkg/agentrun/helpers_test.go index 8fb34bc8e..ef47ad665 100644 --- a/pkg/agentrun/helpers_test.go +++ b/pkg/agentrun/helpers_test.go @@ -281,6 +281,7 @@ func resetRunToPending(t *testing.T, client *pg.Client, runID gid.GID) { WHERE id = $1`, runID.String(), ) + return err }, ) @@ -328,6 +329,7 @@ func bumpRunLeaseGeneration(t *testing.T, client *pg.Client, runID gid.GID) { WHERE id = $1`, runID.String(), ) + return err }, ) diff --git a/pkg/agentrun/worker_test.go b/pkg/agentrun/worker_test.go index 9848d64b6..8a8509f33 100644 --- a/pkg/agentrun/worker_test.go +++ b/pkg/agentrun/worker_test.go @@ -60,6 +60,7 @@ func TestWorker_PicksUpAndCompletes(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) defer cancel() + go func() { _ = runWorker.Run(ctx) }() require.Eventually( @@ -92,6 +93,7 @@ func TestWorker_StopAndResume(t *testing.T) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) { close(toolReady) <-toolRelease + return agent.ToolResult{Content: "work done"}, nil }, ) @@ -122,6 +124,7 @@ func TestWorker_StopAndResume(t *testing.T) { ctx1, cancel1 := context.WithTimeout(context.Background(), 15*time.Second) defer cancel1() + go func() { _ = runWorker.Run(ctx1) }() select { @@ -164,6 +167,7 @@ func TestWorker_StopAndResume(t *testing.T) { ctx2, cancel2 := context.WithTimeout(context.Background(), 15*time.Second) defer cancel2() + go func() { _ = runWorker2.Run(ctx2) }() require.Eventually( @@ -200,6 +204,7 @@ func TestWorker_StopAndResumeAcrossHandoff(t *testing.T) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) { close(toolReady) <-toolRelease + return agent.ToolResult{Content: "child work done"}, nil }, ) @@ -250,6 +255,7 @@ func TestWorker_StopAndResumeAcrossHandoff(t *testing.T) { ctx1, cancel1 := context.WithTimeout(context.Background(), 15*time.Second) defer cancel1() + go func() { _ = runWorker.Run(ctx1) }() select { @@ -295,6 +301,7 @@ func TestWorker_StopAndResumeAcrossHandoff(t *testing.T) { ctx2, cancel2 := context.WithTimeout(context.Background(), 15*time.Second) defer cancel2() + go func() { _ = runWorker2.Run(ctx2) }() require.Eventually( @@ -329,6 +336,7 @@ func TestWorker_StopAndResumeNestedSubAgent(t *testing.T) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) { readyOnce.Do(func() { close(toolReady) }) <-toolRelease + return agent.ToolResult{Content: "inner work done"}, nil }, ) @@ -379,6 +387,7 @@ func TestWorker_StopAndResumeNestedSubAgent(t *testing.T) { ctx1, cancel1 := context.WithTimeout(context.Background(), 15*time.Second) defer cancel1() + go func() { _ = runWorker.Run(ctx1) }() select { @@ -425,6 +434,7 @@ func TestWorker_StopAndResumeNestedSubAgent(t *testing.T) { ctx2, cancel2 := context.WithTimeout(context.Background(), 15*time.Second) defer cancel2() + go func() { _ = runWorker2.Run(ctx2) }() require.Eventually( @@ -459,6 +469,7 @@ func TestWorker_StopAndResumeNestedSubAgentMultiLevel(t *testing.T) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) { readyOnce.Do(func() { close(toolReady) }) <-toolRelease + return agent.ToolResult{Content: "grandchild work done"}, nil }, ) @@ -524,6 +535,7 @@ func TestWorker_StopAndResumeNestedSubAgentMultiLevel(t *testing.T) { ctx1, cancel1 := context.WithTimeout(context.Background(), 15*time.Second) defer cancel1() + go func() { _ = runWorker.Run(ctx1) }() select { @@ -574,6 +586,7 @@ func TestWorker_StopAndResumeNestedSubAgentMultiLevel(t *testing.T) { ctx2, cancel2 := context.WithTimeout(context.Background(), 15*time.Second) defer cancel2() + go func() { _ = runWorker2.Run(ctx2) }() require.Eventually( @@ -697,6 +710,7 @@ func TestWorker_ReclaimedRunDoesNotClobberWinner(t *testing.T) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) { close(toolReady) <-toolRelease + return agent.ToolResult{Content: "work done"}, nil }, ) @@ -735,6 +749,7 @@ func TestWorker_ReclaimedRunDoesNotClobberWinner(t *testing.T) { ctxA, cancelA := context.WithTimeout(context.Background(), 30*time.Second) defer cancelA() + go func() { _ = runWorkerA.Run(ctxA) }() select { @@ -754,6 +769,7 @@ func TestWorker_ReclaimedRunDoesNotClobberWinner(t *testing.T) { ctxB, cancelB := context.WithTimeout(context.Background(), 30*time.Second) defer cancelB() + go func() { _ = runWorkerB.Run(ctxB) }() require.Eventually( @@ -777,6 +793,7 @@ func TestWorker_ReclaimedRunDoesNotClobberWinner(t *testing.T) { func() bool { provider.mu.Lock() defer provider.mu.Unlock() + return provider.calls >= 3 }, 15*time.Second, @@ -815,6 +832,7 @@ func TestWorker_UnknownAgentFails(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) defer cancel() + go func() { _ = runWorker.Run(ctx) }() require.Eventually( @@ -858,6 +876,7 @@ func TestWorker_InvalidInputMessagesFails(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) defer cancel() + go func() { _ = runWorker.Run(ctx) }() require.Eventually( @@ -882,33 +901,49 @@ func TestWorker_SIGTERM(t *testing.T) { return } + // Skip when the test database is unreachable so the parent does not + // wait on a subprocess that skips itself for the same reason and never + // prints READY. + test.PGClient(t) + cmd := exec.Command(os.Args[0], "-test.run=^TestWorker_SIGTERM$") + cmd.Env = append(os.Environ(), "TEST_SIGTERM_SUBPROCESS=1") stdout, err := cmd.StdoutPipe() require.NoError(t, err) + cmd.Stderr = cmd.Stdout require.NoError(t, cmd.Start()) ready := make(chan struct{}) scanDone := make(chan struct{}) - var linesMu sync.Mutex - var lines []string + + var ( + linesMu sync.Mutex + lines []string + ) + snapshotLines := func() string { linesMu.Lock() defer linesMu.Unlock() + return strings.Join(lines, "\n") } + go func() { defer close(scanDone) scanner := bufio.NewScanner(stdout) for scanner.Scan() { line := scanner.Text() + linesMu.Lock() + lines = append(lines, line) linesMu.Unlock() + if line == "READY" { close(ready) } @@ -919,13 +954,16 @@ func TestWorker_SIGTERM(t *testing.T) { case <-ready: case <-time.After(20 * time.Second): _ = cmd.Process.Kill() + t.Fatalf("subprocess did not become ready for SIGTERM\n%s", snapshotLines()) } require.NoError(t, cmd.Process.Signal(syscall.SIGTERM)) + if err := cmd.Wait(); err != nil { t.Fatalf("subprocess failed: %v\n%s", err, snapshotLines()) } + <-scanDone } @@ -967,7 +1005,7 @@ func runSIGTERMSubprocess(t *testing.T) { t.Fatal("tool did not start before SIGTERM") } - fmt.Fprintln(os.Stdout, "READY") + _, _ = fmt.Fprintln(os.Stdout, "READY") select { case <-runWorker.ShutdownBroadcast(): @@ -995,6 +1033,7 @@ func makeBattleTools(workStarted chan<- struct{}) []agent.Tool { func(ctx context.Context, _ workInput) (agent.ToolResult, error) { close(workStarted) <-ctx.Done() + return agent.ToolResult{Content: "interrupted"}, ctx.Err() }, ),