diff --git a/CLAUDE.md b/CLAUDE.md index 0119f5c..ed7d126 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -99,7 +99,8 @@ One `package main`, one file per concern. | `workflow_store.go` | Three-scope workflow persistence: user (ask.json) + repo (`/.ask/workflows/*.json`, committed) + global (`~/.config/ask/workflows/*.json`, machine-local, visible from every project); merged global-first listing (personal-wins), ambiguity-strict name resolution, dir sync on save, cross-scope copy. | | `workflows_screen.go` | Workflows builder screen — list/steps/step editor levels with multi-line prompt textarea. `e` on a selected workflow opens that same textarea (workflow-scoped `promptTarget=="description"`) to edit the workflow's free-text Description; it commits to `workflowDef.Description` and shows as the steps-pane subtitle. | | `workflows_picker.go` | Small centred modal popped on `f` (issues) / `Ctrl+F` (chat) to pick which workflow to run. | -| `pkg/workflow/compile.go` | `CompileWorkflow` — a `Def` becomes an ADK graph: `AgentNode` per step, `loopagent`+`exitlooptool` per loop, `IncludeContentsNone` + `InstructionProvider` per step agent, per-node `RetryConfig`. Returns `Compiled` with the agent-name → step-index map the progress adapter needs. | +| `pkg/workflow/compile.go` | `CompileWorkflow` — a `Def` becomes an ADK graph: `AgentNode` per step, `loopagent`+`exitlooptool` (tail step only) per loop, `IncludeContentsNone` + `InstructionProvider` per step agent, per-node `RetryConfig`. Hands each step a `StepRole{InLoop, IsTail, IsFinal}` so the `ToolsBuilder` attaches position-dependent tools. Returns `Compiled` with the agent-name → step-index map the progress adapter needs. | +| `pkg/tools/artifact.go` | `SaveArtifactTool` (writes through `ctx.Artifacts().Save`; ADK ships only `load_artifacts`) and `WorkflowStepTools(env, isFinal)` — save/load on every step for data handoff, `finish_workflow` on the final step for the user-facing outcome report. | | `pkg/workflow/progress.go` | `Progress` — ADK events → `RunnerListener` callbacks, driven only by events that actually arrived. | | `cmd/ask/workflow_graph.go` | TUI runner: one session for the whole graph, agent swapped for the compiled workflow. | | `pkg/engine/workflow_run.go` | Headless runner + `WorkflowGraphAgent` + `IngestWorkflowMemory`. | @@ -438,14 +439,22 @@ were deleted. Don't add a second one. - A top-level agent step becomes an `AgentNode` wrapping an `llmagent`; nodes are chained `Start -> n0 -> n1 -> …`. - A `kind: "loop"` step becomes an `AgentNode` wrapping ADK's - `loopagent`, whose sub-agents are the inner steps, each carrying - `exitlooptool`. A step breaks the loop by calling `exit_loop` (it sets - `Actions.Escalate`, which is what `loopagent` watches for); otherwise - the loop runs to `MaxIterations`. There is no `decision` argument any - more — loop control is ADK's tool, not an `end_turn` field. + `loopagent`, whose sub-agents are the inner steps. **Only the tail + (last) inner step carries `exitlooptool`**, so only it can break the + loop early: `exit_loop` sets `Actions.Escalate` (what `loopagent` + watches), no other ask tool touches `Actions`, so withholding the tool + from the non-tail steps makes an early break structurally impossible + for them — this replaces the old `end_turn` decision guard, which + caught the violation after the fact and re-prompted. Otherwise the loop + runs to `MaxIterations`. There is no `decision` argument — loop control + is ADK's tool, not an `end_turn` field. - `NodeConfig.RetryConfig` gives per-node retry, replacing the runner's hand-rolled `stepErrorRetry` loop. +`CompileWorkflow` hands each step a `StepRole{InLoop, IsTail, IsFinal}` +via its `ToolsBuilder`, which is what decides the position-dependent +tools: `IsTail` gates `exit_loop`, `IsFinal` gates `finish_workflow`. + Two `llmagent` settings carry the semantics and MUST NOT be dropped: - **`IncludeContents: IncludeContentsNone`** is what isolates a step. @@ -500,8 +509,29 @@ Every step should call `end_turn` once with a `summary` (1-3 sentences); it becomes the step's line in the workflow log via `stepSummaryLine`. A step that ends without it is NOT re-prompted any more — the whole remind/re-prompt machinery is gone — its log line falls back to the first -line of its own output. `finish_workflow` still reports the run's -outcome; the runner reads it from `env.PendingFinishData`. +line of its own output. + +`finish_workflow` reports the run's outcome and the artifacts it produced +(PR links, tickets) to the user. It is attached to the **final step +only** (`WorkflowStepTools(env, isFinal)` in pkg/tools/artifact.go); +`Progress` captures the call straight from the event stream +(progress.go), and the TUI also reads `env.PendingFinishData` as a +backup. Nothing attaches it otherwise — an earlier regression left it +attached to no step at all once the graph became one session. + +### Passing data between steps (artifacts) + +A node's text output is the implicit handoff to the next node. For +structured data — a plan, a diff, notes — a step calls the native +`save_artifact` tool; a later step loads it by name with ADK's +`load_artifacts`. Both are attached to every workflow step by +`WorkflowStepTools`. This works because the whole graph runs as ONE +runner invocation, so the runner's `ArtifactService` (set in +`RunnerBuilder`) spans every step — a save in step 1 is visible to a load +in step 3. ADK ships `load_artifacts` but no save tool; `SaveArtifactTool` +(pkg/tools/artifact.go) is the missing half, writing through +`ctx.Artifacts().Save`. This is the ADK-native replacement for the old +`ask/plans/` notes directories. ### Step instruction assembly @@ -524,9 +554,10 @@ node's input, and `IncludeContentsNone` is what lets the step see it. `ask/plans/` is gone, along with `plans.go`, `clear_plans`, and the `RemindFixPlanDir` re-prompt. Step-to-step handoff is the graph's node -output. Durable "what we learned" goes to `pkg/memory` — `RunWorkflow` -calls `engine.IngestWorkflowMemory` on a clean finish, which is what the -notes directories were badly approximating. +output plus ADK artifacts (see above). Durable "what we learned" goes to +`pkg/memory` — `RunWorkflow` calls `engine.IngestWorkflowMemory` on a +clean finish, which is what the notes directories were badly +approximating. ### Builder screen (`Ctrl+W` / `/workflows`) diff --git a/cmd/ask/workflow_graph.go b/cmd/ask/workflow_graph.go index a3e9066..a6443b5 100644 --- a/cmd/ask/workflow_graph.go +++ b/cmd/ask/workflow_graph.go @@ -7,6 +7,7 @@ import ( "github.com/Cidan/ask/pkg/engine" "github.com/Cidan/ask/pkg/providers" + "github.com/Cidan/ask/pkg/tools" "github.com/Cidan/ask/pkg/workflow" adkmodel "google.golang.org/adk/v2/model" adktool "google.golang.org/adk/v2/tool" @@ -127,10 +128,12 @@ func tuiWorkflowCompileConfig(sess *agentSession, def workflow.Def, src workflow ModelBuilder: func(ctx context.Context, step workflow.Step) (adkmodel.LLM, error) { return workflowStepModel(ctx, sess, step) }, - ToolsBuilder: func(ctx context.Context, step workflow.Step, inLoop bool) ([]adktool.Tool, error) { - return engine.AsADKTools(sess.currentTools()) + ToolsBuilder: func(ctx context.Context, step workflow.Step, role workflow.StepRole) ([]adktool.Tool, error) { + list := append([]tools.Tool(nil), sess.currentTools()...) + list = append(list, tools.WorkflowStepTools(sess.env, role.IsFinal)...) + return engine.AsADKTools(list) }, - ToolsetsBuilder: func(ctx context.Context, step workflow.Step, inLoop bool) ([]adktool.Toolset, error) { + ToolsetsBuilder: func(ctx context.Context, step workflow.Step, role workflow.StepRole) ([]adktool.Toolset, error) { var toolsets []adktool.Toolset if sess.mcp != nil { toolsets = append(toolsets, sess.mcp.Toolsets()...) diff --git a/pkg/engine/run.go b/pkg/engine/run.go index 5b66185..0247ed6 100644 --- a/pkg/engine/run.go +++ b/pkg/engine/run.go @@ -105,6 +105,13 @@ type ToolFactoryArgs struct { EventListener EventListener InteractionHandler InteractionHandler AttachWebSearch bool + // WorkflowStep attaches the workflow-step tools (save_artifact, + // load_artifacts) so a step can pass data to a later one. + WorkflowStep bool + // WorkflowFinalStep also attaches finish_workflow, which reports the + // run's outcome and created artifacts to the user. Only meaningful + // with WorkflowStep. + WorkflowFinalStep bool } // ToolFactory builds a slice of Tools for an engine turn. diff --git a/pkg/engine/workflow_artifact_integration_test.go b/pkg/engine/workflow_artifact_integration_test.go new file mode 100644 index 0000000..49c8dcb --- /dev/null +++ b/pkg/engine/workflow_artifact_integration_test.go @@ -0,0 +1,177 @@ +package engine_test + +import ( + "context" + "iter" + "strings" + "sync" + "testing" + + "github.com/Cidan/ask/pkg/config" + "github.com/Cidan/ask/pkg/engine" + "github.com/Cidan/ask/pkg/providers" + // Imported for its init(): registers the tool factory, so workflow + // steps get the real core toolset plus save_artifact / load_artifacts. + _ "github.com/Cidan/ask/pkg/tools" + "github.com/Cidan/ask/pkg/workflow" + adkmodel "google.golang.org/adk/v2/model" + "google.golang.org/genai" +) + +// scriptedModel drives one step at a time by inspecting the marker in the +// step's system instruction and counting calls per marker. It is a full +// model.LLM, not the internal test mock, because this test lives in the +// external engine_test package to avoid the tools -> engine import cycle. +type scriptedModel struct { + mu sync.Mutex + calls map[string]int + // readerHandoff is what the reader saw on its FIRST call — the step + // handoff, before any artifact was loaded. + readerHandoff string + // readerContent captures every request text the reader step saw on + // its second call, so the test can assert the artifact crossed over. + readerContent []string +} + +func (m *scriptedModel) Name() string { return "scripted" } + +func (m *scriptedModel) GenerateContent(_ context.Context, req *adkmodel.LLMRequest, _ bool) iter.Seq2[*adkmodel.LLMResponse, error] { + m.mu.Lock() + defer m.mu.Unlock() + + sys := systemText(req) + switch { + case strings.Contains(sys, "SAVE_STEP_MARKER"): + n := m.calls["save"] + m.calls["save"]++ + if n == 0 { + return one(functionCall("save_artifact", map[string]any{ + "name": "plan.md", + "content": "SECRET_PLAN_BODY", + "description": "save the plan", + })) + } + return one(finalText("saved the plan")) + + case strings.Contains(sys, "READ_STEP_MARKER"): + n := m.calls["read"] + m.calls["read"]++ + if n == 0 { + m.readerHandoff = requestText(req) + return one(functionCall("load_artifacts", map[string]any{ + "artifact_names": []any{"plan.md"}, + })) + } + // Second call: load_artifacts' ProcessRequest has injected the + // artifact content into the request by now. + m.readerContent = append(m.readerContent, requestText(req)) + return one(finalText("read the plan")) + } + return one(finalText("done")) +} + +func systemText(req *adkmodel.LLMRequest) string { + if req == nil || req.Config == nil || req.Config.SystemInstruction == nil { + return "" + } + var b strings.Builder + for _, p := range req.Config.SystemInstruction.Parts { + if p != nil { + b.WriteString(p.Text) + } + } + return b.String() +} + +func requestText(req *adkmodel.LLMRequest) string { + var b strings.Builder + for _, c := range req.Contents { + if c == nil { + continue + } + for _, p := range c.Parts { + if p != nil && p.Text != "" { + b.WriteString(p.Text) + b.WriteByte('\n') + } + } + } + return b.String() +} + +func functionCall(name string, args map[string]any) *adkmodel.LLMResponse { + return &adkmodel.LLMResponse{ + Content: &genai.Content{ + Role: genai.RoleModel, + Parts: []*genai.Part{genai.NewPartFromFunctionCall(name, args)}, + }, + } +} + +func finalText(text string) *adkmodel.LLMResponse { + return &adkmodel.LLMResponse{ + Content: &genai.Content{Role: genai.RoleModel, Parts: []*genai.Part{genai.NewPartFromText(text)}}, + FinishReason: genai.FinishReasonStop, + } +} + +func one(resp *adkmodel.LLMResponse) iter.Seq2[*adkmodel.LLMResponse, error] { + return func(yield func(*adkmodel.LLMResponse, error) bool) { yield(resp, nil) } +} + +// TestWorkflowArtifactHandoff is the full round trip: a two-step workflow +// on the real ADK graph, where step 1 saves an artifact and step 2 loads +// it. It proves the wiring end to end — the tool factory attaches +// save_artifact/load_artifacts to every step, the graph runs as one +// runner invocation so the ArtifactService spans both steps, and the +// content saved by step 1 reaches step 2's model. +func TestWorkflowArtifactHandoff(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + tmp := t.TempDir() + + scripted := &scriptedModel{calls: map[string]int{}} + origBuilder := engine.ModelBuilder + engine.ModelBuilder = func(_ context.Context, _ *providers.AgentProviderSpec, _ config.Config, _ string) (adkmodel.LLM, error) { + return scripted, nil + } + defer func() { engine.ModelBuilder = origBuilder }() + + eng := engine.New(engine.Options{ + Config: config.Config{Provider: "vertex"}, + InteractionHandler: engine.HeadlessInteractionHandler{AutoApproveTools: true}, + }) + + def := workflow.Def{ + Name: "artifact-handoff", + Steps: []workflow.Step{ + {Name: "saver", Prompt: "SAVE_STEP_MARKER: save the plan as an artifact."}, + {Name: "reader", Prompt: "READ_STEP_MARKER: load the plan artifact and use it."}, + }, + } + src := workflow.NewTextSource(1, "Artifact Handoff Source") + + if err := eng.RunWorkflow(context.Background(), tmp, 1, def, src); err != nil { + t.Fatalf("RunWorkflow: %v", err) + } + + scripted.mu.Lock() + defer scripted.mu.Unlock() + + if scripted.calls["save"] == 0 { + t.Fatal("the saver step never ran") + } + if len(scripted.readerContent) == 0 { + t.Fatal("the reader step never reached its second call, so load_artifacts never resolved") + } + // The secret must NOT arrive through the normal step handoff — the + // saver's node output is "saved the plan", and IncludeContentsNone + // keeps the saver's tool calls out of the reader's context. So its + // only path to the reader is load_artifacts. + if strings.Contains(scripted.readerHandoff, "SECRET_PLAN_BODY") { + t.Fatalf("the secret leaked through the step handoff, not the artifact:\n%s", scripted.readerHandoff) + } + joined := strings.Join(scripted.readerContent, "\n") + if !strings.Contains(joined, "SECRET_PLAN_BODY") { + t.Errorf("the artifact saved by step 1 did not reach step 2's model via load_artifacts:\n%s", joined) + } +} diff --git a/pkg/engine/workflow_run.go b/pkg/engine/workflow_run.go index 5add45d..adc4a2c 100644 --- a/pkg/engine/workflow_run.go +++ b/pkg/engine/workflow_run.go @@ -27,7 +27,7 @@ func WorkflowCompileConfig(e *Engine, cwd string, tabID int, def workflow.Def, s ModelBuilder: func(ctx context.Context, step workflow.Step) (adkmodel.LLM, error) { return buildStepModel(ctx, e, step) }, - ToolsBuilder: func(ctx context.Context, step workflow.Step, inLoop bool) ([]tool.Tool, error) { + ToolsBuilder: func(ctx context.Context, step workflow.Step, role workflow.StepRole) ([]tool.Tool, error) { var agentTools []Tool if tf := GetDefaultToolFactory(); tf != nil { agentTools = tf(ToolFactoryArgs{ @@ -37,11 +37,13 @@ func WorkflowCompileConfig(e *Engine, cwd string, tabID int, def workflow.Def, s EventListener: e.opts.EventListener, InteractionHandler: e.opts.InteractionHandler, AttachWebSearch: true, + WorkflowStep: true, + WorkflowFinalStep: role.IsFinal, }) } return AsADKTools(agentTools) }, - ToolsetsBuilder: func(ctx context.Context, step workflow.Step, inLoop bool) ([]tool.Toolset, error) { + ToolsetsBuilder: func(ctx context.Context, step workflow.Step, role workflow.StepRole) ([]tool.Toolset, error) { var toolsets []tool.Toolset if skillTS, err := NewSkillToolset(ctx, cwd); err == nil && skillTS != nil { toolsets = append(toolsets, skillTS) diff --git a/pkg/tools/artifact.go b/pkg/tools/artifact.go new file mode 100644 index 0000000..f469034 --- /dev/null +++ b/pkg/tools/artifact.go @@ -0,0 +1,68 @@ +package tools + +import ( + "errors" + "strings" + + "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/tool/loadartifactstool" + "google.golang.org/genai" +) + +// SaveArtifactParams is the save_artifact tool's input. +type SaveArtifactParams struct { + Name string `json:"name" jsonschema:"artifact name, e.g. plan.md — later steps load it by this name"` + Content string `json:"content" jsonschema:"the artifact's text content"` + Description string `json:"description" jsonschema:"one short human-readable phrase (under 10 words) telling the user what this call is doing"` +} + +// SaveArtifactResult is the save_artifact tool's response. +type SaveArtifactResult struct { + Name string `json:"name,omitempty"` + Version int64 `json:"version,omitempty" jsonschema:"the saved version number"` +} + +// SaveArtifactTool lets a workflow step hand structured data to a later +// step. The artifact lives in the run's ArtifactService, which spans the +// whole graph because the graph runs as one runner invocation, so a step +// saves plan.md and a downstream step reads it back with load_artifacts. +// +// ADK ships load_artifacts but no save tool; this is the missing half. +func SaveArtifactTool() Tool { + return NewTypedTool( + "save_artifact", + "Save a named artifact (a plan, a diff, notes) so a later workflow step can load it by name. Use this to hand structured output forward rather than restating it.", + func(ctx agent.Context, p SaveArtifactParams) (SaveArtifactResult, error) { + name := strings.TrimSpace(p.Name) + if name == "" { + return SaveArtifactResult{}, errors.New("name is required") + } + arts := ctx.Artifacts() + if arts == nil { + return SaveArtifactResult{}, errors.New("artifacts are not available in this context") + } + resp, err := arts.Save(ctx, name, genai.NewPartFromText(p.Content)) + if err != nil { + return SaveArtifactResult{}, err + } + var version int64 + if resp != nil { + version = resp.Version + } + return SaveArtifactResult{Name: name, Version: version}, nil + }, + ) +} + +// WorkflowStepTools returns the position-dependent tools every workflow +// step needs: save_artifact and load_artifacts on every step so steps can +// pass data forward, plus finish_workflow on the final step only, which +// reports the run's outcome and created artifacts (PR links, tickets) to +// the user. +func WorkflowStepTools(env *ToolEnv, isFinal bool) []Tool { + out := []Tool{SaveArtifactTool(), loadartifactstool.New()} + if isFinal { + out = append(out, FinishWorkflowTool(env)) + } + return out +} diff --git a/pkg/tools/artifact_test.go b/pkg/tools/artifact_test.go new file mode 100644 index 0000000..2780a4e --- /dev/null +++ b/pkg/tools/artifact_test.go @@ -0,0 +1,115 @@ +package tools + +import ( + "context" + "strings" + "testing" + + "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/artifact" + "google.golang.org/genai" +) + +// A non-final step can pass data forward (save/load) but must not be able +// to report the run's outcome; only the final step gets finish_workflow. +func TestWorkflowStepTools(t *testing.T) { + env := NewToolEnv(t.TempDir(), 1, true, false, nil, nil) + + names := func(ts []Tool) map[string]bool { + m := map[string]bool{} + for _, tl := range ts { + m[tl.Name()] = true + } + return m + } + + mid := names(WorkflowStepTools(env, false)) + if !mid["save_artifact"] || !mid["load_artifacts"] { + t.Errorf("every workflow step needs save/load, got %v", mid) + } + if mid["finish_workflow"] { + t.Error("a non-final step must not get finish_workflow") + } + + final := names(WorkflowStepTools(env, true)) + if !final["save_artifact"] || !final["load_artifacts"] || !final["finish_workflow"] { + t.Errorf("the final step needs save/load/finish, got %v", final) + } +} + +// save_artifact reports a real error (not a silent field) when the +// context has no artifact service, so the failure reaches the model. +func TestSaveArtifactTool_NoArtifactService(t *testing.T) { + _, err := runTypedTool[SaveArtifactResult](t, SaveArtifactTool(), + SaveArtifactParams{Name: "plan.md", Content: "hi", Description: "save"}) + if err == nil || !strings.Contains(err.Error(), "artifacts are not available") { + t.Errorf("save without an artifact service should error, got %v", err) + } +} + +func TestSaveArtifactTool_RequiresName(t *testing.T) { + _, err := runTypedTool[SaveArtifactResult](t, SaveArtifactTool(), + SaveArtifactParams{Name: " ", Content: "hi", Description: "save"}) + if err == nil || !strings.Contains(err.Error(), "name is required") { + t.Errorf("save without a name should error, got %v", err) + } +} + +// fakeArtifacts is a minimal agent.Artifacts backed by a map, so a test +// can watch what save_artifact writes without standing up a runner. +type fakeArtifacts struct{ store map[string]*genai.Part } + +func (f *fakeArtifacts) Save(_ context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { + f.store[name] = data + return &artifact.SaveResponse{Version: int64(len(f.store))}, nil +} +func (f *fakeArtifacts) Load(_ context.Context, name string) (*artifact.LoadResponse, error) { + return &artifact.LoadResponse{Part: f.store[name]}, nil +} +func (f *fakeArtifacts) LoadVersion(_ context.Context, name string, _ int) (*artifact.LoadResponse, error) { + return &artifact.LoadResponse{Part: f.store[name]}, nil +} +func (f *fakeArtifacts) List(context.Context) (*artifact.ListResponse, error) { + names := make([]string, 0, len(f.store)) + for n := range f.store { + names = append(names, n) + } + return &artifact.ListResponse{FileNames: names}, nil +} + +// artifactsCtx overrides just Artifacts() on a real agent.Context. +type artifactsCtx struct { + agent.Context + arts agent.Artifacts +} + +func (a artifactsCtx) Artifacts() agent.Artifacts { return a.arts } + +// save_artifact writes the content to the run's artifact store under the +// given name — the handoff a later step reads back with load_artifacts. +func TestSaveArtifactTool_WritesToArtifactStore(t *testing.T) { + arts := &fakeArtifacts{store: map[string]*genai.Part{}} + ctx := artifactsCtx{Context: testAgentCtx(), arts: arts} + + runner, ok := SaveArtifactTool().(interface { + Run(agent.Context, any) (map[string]any, error) + }) + if !ok { + t.Fatal("save_artifact does not implement the ADK Run contract") + } + res, err := runner.Run(ctx, map[string]any{ + "name": "plan.md", + "content": "the plan body", + "description": "save the plan", + }) + if err != nil { + t.Fatalf("save_artifact: %v", err) + } + if res["name"] != "plan.md" { + t.Errorf("result name = %v, want plan.md", res["name"]) + } + saved := arts.store["plan.md"] + if saved == nil || saved.Text != "the plan body" { + t.Errorf("artifact store did not receive the content: %+v", saved) + } +} diff --git a/pkg/tools/core.go b/pkg/tools/core.go index a06c968..9210435 100644 --- a/pkg/tools/core.go +++ b/pkg/tools/core.go @@ -28,7 +28,11 @@ func BuildCoreTools(args engine.ToolFactoryArgs, attachWebSearch bool) []Tool { registryFunc := func() []Tool { return nil } - return CoreTools(env, registryFunc, attachWebSearch) + core := CoreTools(env, registryFunc, attachWebSearch) + if args.WorkflowStep { + core = append(core, WorkflowStepTools(env, args.WorkflowFinalStep)...) + } + return core } // BuildSubagentTools constructs the centralized toolset configured for an isolated subagent session. diff --git a/pkg/workflow/compile.go b/pkg/workflow/compile.go index 561f423..73554da 100644 --- a/pkg/workflow/compile.go +++ b/pkg/workflow/compile.go @@ -29,8 +29,10 @@ type WorkflowAgentConfig struct { // ModelBuilder resolves the LLM for one step. Required. ModelBuilder func(ctx context.Context, step Step) (model.LLM, error) // ToolsBuilder and ToolsetsBuilder supply the step's tool surface. - ToolsBuilder func(ctx context.Context, step Step, inLoop bool) ([]tool.Tool, error) - ToolsetsBuilder func(ctx context.Context, step Step, inLoop bool) ([]tool.Toolset, error) + // StepRole tells them where the step sits, so the builder can attach + // finish_workflow to the final step only, and so on. + ToolsBuilder func(ctx context.Context, step Step, role StepRole) ([]tool.Tool, error) + ToolsetsBuilder func(ctx context.Context, step Step, role StepRole) ([]tool.Toolset, error) // InstructionBuilder renders the step's system instruction. Defaults // to BuildStepInstruction. InstructionBuilder func(step Step, src Source, pc *StepPromptCtx) string @@ -44,6 +46,20 @@ type WorkflowAgentConfig struct { // RetryConfig, which also handles the backoff. const workflowDefaultMaxRetries = 3 +// StepRole tells a ToolsBuilder where a step sits in the workflow, so it +// can decide which position-dependent tools to attach. +type StepRole struct { + // InLoop is true when the step runs inside a loop container. + InLoop bool + // IsTail is true when the step is the last inner step of its loop — + // the only step allowed to break the loop early via exit_loop. + IsTail bool + // IsFinal is true when the step is the last thing the whole workflow + // runs (the last top-level step, or the tail of a final loop). Only + // this step gets finish_workflow. + IsFinal bool +} + // Compiled is a Def rendered as an executable ADK graph, plus the lookup // the event adapter needs to attribute events back to steps. type Compiled struct { @@ -186,16 +202,24 @@ func buildStepAgent(ctx context.Context, cfg WorkflowAgentConfig, step Step, ste return nil, fmt.Errorf("workflow compile: model for step %q: %w", step.Name, err) } - inLoop := loop != nil + role := StepRole{InLoop: loop != nil, IsFinal: isFinal} + if loop != nil { + role.IsTail = loop.IsTail + } var tools []tool.Tool if cfg.ToolsBuilder != nil { - tools, err = cfg.ToolsBuilder(ctx, step, inLoop) + tools, err = cfg.ToolsBuilder(ctx, step, role) if err != nil { return nil, fmt.Errorf("workflow compile: tools for step %q: %w", step.Name, err) } } - if inLoop { + // Only the tail step of a loop may break out early. exit_loop sets + // Actions.Escalate, which is the only way an ask tool ends a + // loopagent, and no other tool touches Actions — so withholding it + // from the non-tail steps makes an early break structurally + // impossible for them, rather than something caught after the fact. + if role.InLoop && role.IsTail { tools, err = withExitLoopTool(tools) if err != nil { return nil, fmt.Errorf("workflow compile: step %q: %w", step.Name, err) @@ -204,7 +228,7 @@ func buildStepAgent(ctx context.Context, cfg WorkflowAgentConfig, step Step, ste var toolsets []tool.Toolset if cfg.ToolsetsBuilder != nil { - toolsets, err = cfg.ToolsetsBuilder(ctx, step, inLoop) + toolsets, err = cfg.ToolsetsBuilder(ctx, step, role) if err != nil { return nil, fmt.Errorf("workflow compile: toolsets for step %q: %w", step.Name, err) } diff --git a/pkg/workflow/compile_test.go b/pkg/workflow/compile_test.go index 7ad1b0f..cbe1a39 100644 --- a/pkg/workflow/compile_test.go +++ b/pkg/workflow/compile_test.go @@ -231,3 +231,104 @@ func TestBuildStepInstruction_LoopFraming(t *testing.T) { t.Errorf("loop control moved to exit_loop; instruction should not mention a decision arg: %q", got) } } + +// Only the tail step is told how to break, because only the tail step has +// the exit_loop tool. A non-tail step is told it cannot end the loop. +func TestBuildStepInstruction_NonTailLoopStepCannotBreak(t *testing.T) { + src := NewTextSource(1, "src") + step := Step{Name: "edit", Prompt: "Edit."} + loop := &LoopPromptCtx{Name: "fix", MaxIterations: 5, IsTail: false} + + got := BuildStepInstruction(step, src, &StepPromptCtx{Loop: loop}) + if strings.Contains(got, "exit_loop") { + t.Errorf("a non-tail step must not be told about exit_loop: %q", got) + } + if !strings.Contains(got, "NOT its last step") { + t.Errorf("a non-tail step should be told it cannot end the loop: %q", got) + } +} + +// The final step is told to report the run's outcome with finish_workflow. +func TestBuildStepInstruction_FinalStepReportsWithFinishWorkflow(t *testing.T) { + src := NewTextSource(1, "src") + step := Step{Name: "ship", Prompt: "Open the PR."} + + final := BuildStepInstruction(step, src, &StepPromptCtx{IsWorkflowFinalStep: true}) + if !strings.Contains(final, "finish_workflow") { + t.Errorf("the final step must be told to call finish_workflow: %q", final) + } + + mid := BuildStepInstruction(step, src, &StepPromptCtx{}) + if strings.Contains(mid, "finish_workflow") { + t.Errorf("a non-final step must not be told about finish_workflow: %q", mid) + } + // Every step is told it can pass data forward. + if !strings.Contains(mid, "save_artifact") { + t.Errorf("every step should learn about save_artifact: %q", mid) + } +} + +// recordRoles compiles def with a ToolsBuilder that records the StepRole +// handed to each step, keyed by step name. +func recordRoles(t *testing.T, def Def) map[string]StepRole { + t.Helper() + roles := map[string]StepRole{} + cfg := testCompileConfig(def) + cfg.ToolsBuilder = func(_ context.Context, step Step, role StepRole) ([]tool.Tool, error) { + roles[step.Name] = role + return nil, nil + } + if _, err := CompileWorkflow(context.Background(), cfg); err != nil { + t.Fatalf("compile: %v", err) + } + return roles +} + +// The role handed to each step is what decides its position-dependent +// tools: IsTail gates exit_loop (only the tail step may break a loop), +// IsFinal gates finish_workflow (only the last step reports the outcome). +func TestCompileWorkflow_StepRoles(t *testing.T) { + def := Def{Name: "wf", Steps: []Step{ + {Name: "a", Prompt: "a"}, + {Name: "fix", Kind: "loop", Steps: []Step{ + {Name: "b", Prompt: "b"}, + {Name: "c", Prompt: "c"}, + {Name: "d", Prompt: "d"}, + }}, + {Name: "e", Prompt: "e"}, + }} + roles := recordRoles(t, def) + + want := map[string]StepRole{ + "a": {InLoop: false, IsTail: false, IsFinal: false}, + "b": {InLoop: true, IsTail: false, IsFinal: false}, + "c": {InLoop: true, IsTail: false, IsFinal: false}, + "d": {InLoop: true, IsTail: true, IsFinal: false}, // tail, but the loop is not the final step + "e": {InLoop: false, IsTail: false, IsFinal: true}, + } + for name, w := range want { + if roles[name] != w { + t.Errorf("role[%q] = %+v, want %+v", name, roles[name], w) + } + } +} + +// When the workflow ends in a loop, the loop's tail step is BOTH the +// break point and the final step, so it gets exit_loop and finish_workflow. +func TestCompileWorkflow_FinalLoopTailIsFinal(t *testing.T) { + def := Def{Name: "wf", Steps: []Step{ + {Name: "a", Prompt: "a"}, + {Name: "fix", Kind: "loop", Steps: []Step{ + {Name: "b", Prompt: "b"}, + {Name: "c", Prompt: "c"}, + }}, + }} + roles := recordRoles(t, def) + + if r := roles["b"]; r.IsTail || r.IsFinal { + t.Errorf("non-tail loop step b = %+v, want neither tail nor final", r) + } + if r := roles["c"]; !r.IsTail || !r.IsFinal { + t.Errorf("tail of the final loop c = %+v, want both tail and final", r) + } +} diff --git a/pkg/workflow/runner.go b/pkg/workflow/runner.go index a86bfad..95e9899 100644 --- a/pkg/workflow/runner.go +++ b/pkg/workflow/runner.go @@ -76,8 +76,17 @@ func BuildStepInstruction(step Step, source Source, pc *StepPromptCtx) string { b.WriteString(ref) } var loop *LoopPromptCtx + isFinal := false if pc != nil { loop = pc.Loop + isFinal = pc.IsWorkflowFinalStep + } + b.WriteString("\n\nTo hand structured output to a later step — a plan, a diff, notes — call save_artifact " + + "with a name, and the later step loads it with load_artifacts. Prefer this over restating large output in your summary.") + if isFinal { + b.WriteString(" You are the FINAL step of this workflow: before you finish, call finish_workflow with a " + + "description of the outcome and the list of artifacts it produced — every PR, issue, or link the user " + + "needs. This is how the user learns what the run created.") } b.WriteString("\n\n") b.WriteString(EndTurnInstructionBlock(loop)) @@ -103,12 +112,16 @@ func EndTurnInstructionBlock(loop *LoopPromptCtx) string { "a `summary` of 1-3 sentences describing what you did and the outcome. This records your progress in the " + "workflow log; it does not cut your turn short.") if loop != nil { - b.WriteString(" You are running inside a loop: when the loop's exit goal above is met, call the " + - "exit_loop tool to end the loop. If it is not met, do not call exit_loop — the loop advances to its " + - "next iteration on its own, and stops by itself after the iteration limit.") - if !loop.IsTail { - b.WriteString(" Later steps in this iteration still have work to do, so only call exit_loop if the " + - "goal is already fully met.") + if loop.IsTail { + // Only the tail step is given the exit_loop tool, so only the + // tail step is told how to break. + b.WriteString(" You are the last step of this loop iteration: when the loop's exit goal above is met, " + + "call the exit_loop tool to end the loop. If it is not met, do not call exit_loop — the loop advances " + + "to its next iteration on its own, and stops by itself after the iteration limit.") + } else { + b.WriteString(" You are running inside a loop but you are NOT its last step, so you cannot end the loop — " + + "only the final step of the iteration decides whether to continue or stop. Do your part and end your " + + "turn.") } } return b.String()