diff --git a/cmd/ask/agent_provider.go b/cmd/ask/agent_provider.go index 7b1c9fe..218e1cd 100644 --- a/cmd/ask/agent_provider.go +++ b/cmd/ask/agent_provider.go @@ -172,6 +172,7 @@ func setupAgentSessionTools(s *agentSession, cfg askConfig) { agentEndTurnTool(env), agentSearchToolsTool(s.deferredTools), agentInvokeToolTool(s.deferredTools, s.isCoreToolName, env), + agentLoadArtifactsTool(), } if !s.args.InWorkflow { s.coreTools = append(s.coreTools, agentFinalizedPlanTool(env)) @@ -183,7 +184,6 @@ func setupAgentSessionTools(s *agentSession, cfg askConfig) { s.coreTools = append(s.coreTools, agentWorkflowTools(env)...) } s.coreTools = append(s.coreTools, agentWebSearchTool(env)) - s.coreTools = wrapFileToolsWithMemory(s.coreTools, s.args.Cwd) s.coreTools = wrapContextAwareTools(s.coreTools, s.args.Cwd, discoverRules(s.args.Cwd)) s.deferredBase = agentLinearTools(env) s.deferredBase = append(s.deferredBase, agentMemoryIndexTool(env)) diff --git a/cmd/ask/aliases.go b/cmd/ask/aliases.go index 8e79792..02e69ce 100644 --- a/cmd/ask/aliases.go +++ b/cmd/ask/aliases.go @@ -189,6 +189,7 @@ var ( agentWebSearchTool = tools.WebSearchTool agentLoadMemoryTool = tools.LoadMemoryTool agentPreloadMemoryTool = tools.PreloadMemoryTool + agentLoadArtifactsTool = tools.LoadArtifactsTool ) const ( @@ -238,10 +239,6 @@ func agentMemorySystemBlock(cwd string) string { return memory.SystemBlock(context.Background(), cwd) } -func wrapFileToolsWithMemory(ts []tools.Tool, cwd string) []tools.Tool { - return tools.WrapFileToolsWithMemory(ts, cwd) -} - func agentMemoryIndexTool(env *agentToolEnv) tools.Tool { return tools.MemoryIndexTool(env.Cwd, env.RequestApproval) } @@ -347,3 +344,4 @@ func buildAgentSystemPrompt(args ProviderSessionArgs) string { GitStatusFn: agentGitStatus, }) } +} diff --git a/cmd/ask/coordinator.go b/cmd/ask/coordinator.go index bc1f2a2..54fc5af 100644 --- a/cmd/ask/coordinator.go +++ b/cmd/ask/coordinator.go @@ -2,14 +2,15 @@ package main import ( "context" - "errors" "fmt" - "strings" "sync" tea "charm.land/bubbletea/v2" "github.com/Cidan/ask/pkg/engine" + "github.com/Cidan/ask/pkg/providers" "github.com/Cidan/ask/pkg/workflow" + adkmodel "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/tool" ) // Coordinator manages the background execution of all in-process agent sessions @@ -281,96 +282,6 @@ func (l tuiWorkflowListener) OnNote(tabID int, text string) { } // ExecuteStep implements workflow.StepExecutor for Coordinator. -func (c *Coordinator) ExecuteStep(ctx context.Context, cwd string, tabID int, step workflow.Step, prompt string, isFinal bool) (workflow.StepResult, error) { - prov := providerByID(step.Provider) - if prov == nil { - return workflow.StepResult{}, fmt.Errorf("provider not registered: %s", step.Provider) - } - - args := ProviderSessionArgs{ - Cwd: cwd, - TabID: tabID, - Model: step.Model, - Effort: "medium", - SkipAllPermissions: true, - InWorkflow: true, - IsWorkflowFinalStep: isFinal, - } - - proc, ch, err := prov.StartSession(args) - if err != nil { - return workflow.StepResult{}, err - } - - session, ok := proc.payload.(*agentSession) - if !ok { - return workflow.StepResult{}, errors.New("proc payload is not an agent session") - } - c.SetSession(tabID, session) - - err = session.queueTurn(prompt) - if err != nil { - session.shutdown() - c.RemoveSession(tabID) - return workflow.StepResult{}, err - } - - var stepResult string - var stepErr error -stepLoop: - for msg := range ch { - switch m := msg.(type) { - case assistantTextMsg: - stepResult += m.text - case providerDoneMsg: - if m.err != nil { - stepErr = m.err - } else if m.res.IsError { - stepErr = fmt.Errorf("step failed: %s", m.res.Result) - } else { - stepResult = m.res.Result - } - case turnCompleteMsg: - break stepLoop - } - } - - session.shutdown() - c.RemoveSession(tabID) - - if stepErr != nil { - return workflow.StepResult{}, stepErr - } - - summary := "" - decision := "" - if session.env.PendingEndTurn != nil { - summary = session.env.PendingEndTurn.Summary - decision = session.env.PendingEndTurn.Decision - } - if summary == "" && strings.TrimSpace(stepResult) != "" { - firstLine := strings.TrimSpace(strings.Split(strings.TrimSpace(stepResult), "\n")[0]) - if len(firstLine) > 200 { - firstLine = firstLine[:200] + "…" - } - summary = firstLine - } - - var finishData *workflow.FinishData - if session.env.PendingFinishData != nil { - finishData = &workflow.FinishData{ - Description: session.env.PendingFinishData.Description, - Artifacts: session.env.PendingFinishData.Artifacts, - } - } - - return workflow.StepResult{ - Output: stepResult, - Summary: summary, - Decision: decision, - FinishData: finishData, - }, nil -} // RunWorkflow executes a workflow synchronously step by step in the background. func (c *Coordinator) RunWorkflow(ctx context.Context, tabID int, def workflowDef, src workflowSource) (finalizedPlanReply, error) { @@ -407,8 +318,59 @@ func (c *Coordinator) RunWorkflow(ctx context.Context, tabID int, def workflowDe } listener := tuiWorkflowListener{tabID: tabID} - runner := workflow.NewRunner(workflow.GlobalTracker(), c, listener) - runState, err := runner.Run(ctx, rootCwd, tabID, toPkgWorkflowDef(def), src) + + cfg := workflow.WorkflowAgentConfig{ + Def: toPkgWorkflowDef(def), + Source: src, + Cwd: rootCwd, + TabID: tabID, + ModelBuilder: func(ctx context.Context, step workflow.Step) (adkmodel.LLM, error) { + providerID := step.Provider + if providerID == "" { + providerID = "vertex" + } + spec, ok := providers.GetAgentProviderSpec(providerID) + if !ok || spec == nil { + return nil, fmt.Errorf("unknown provider %q", providerID) + } + config, _ := loadConfig() + modelID := providers.CanonicalVertexModelID(step.Model, spec.DefaultModel) + return engine.ModelBuilder(ctx, spec, toPkgConfig(config), modelID) + }, + ToolsBuilder: func(ctx context.Context, step workflow.Step, isLoop bool) ([]tool.Tool, error) { + var agentTools []engine.Tool + if tf := engine.GetDefaultToolFactory(); tf != nil { + agentTools = tf(engine.ToolFactoryArgs{ + Cwd: rootCwd, + TabID: tabID, + SkipPermissions: true, + AttachWebSearch: true, + }) + } + return engine.AsADKTools(agentTools) + }, + ToolsetsBuilder: func(ctx context.Context, step workflow.Step, isLoop bool) ([]tool.Toolset, error) { + var toolsets []tool.Toolset + if skillTS, err := engine.NewSkillToolset(ctx, rootCwd); err == nil && skillTS != nil { + toolsets = append(toolsets, skillTS) + } + return toolsets, nil + }, + InstructionBuilder: func(step workflow.Step, isStart bool, isFinal bool, loopCtx *workflow.LoopPromptCtx, notesDir, prevNotesDir string) string { + pc := &workflow.StepPromptCtx{ + Loop: loopCtx, + NotesDir: notesDir, + PrevNotesDir: prevNotesDir, + IsStartStep: isStart, + IsWorkflowFinalStep: isFinal, + } + return workflow.BuildStepPrompt(step, src, nil, pc) + }, + SessionService: engine.NewFileSessionService("ask-workflow", rootCwd), + } + + runner := workflow.NewRunner(workflow.GlobalTracker(), cfg) + runState, err := runner.Run(ctx, listener) if err != nil { return finalizedPlanReply{}, err } diff --git a/pkg/engine/engine.go b/pkg/engine/engine.go index fa74d80..06ee1fa 100644 --- a/pkg/engine/engine.go +++ b/pkg/engine/engine.go @@ -54,11 +54,12 @@ func (e *Engine) SystemPrompt(cwd string, inWorkflow bool) string { // BuildWorkflowAgent constructs an ADK agent hierarchy (sequentialagent, loopagent, exitlooptool) // for the given workflow definition using the engine's model and tool configuration. -func (e *Engine) BuildWorkflowAgent(ctx context.Context, cwd string, def workflow.Def, src workflow.Source) (agent.Agent, error) { - cfg := workflow.WorkflowAgentConfig{ +func (e *Engine) BuildWorkflowAgentConfig(ctx context.Context, cwd string, tabID int, def workflow.Def, src workflow.Source) workflow.WorkflowAgentConfig { + return workflow.WorkflowAgentConfig{ Def: def, Source: src, Cwd: cwd, + TabID: tabID, ModelBuilder: func(ctx context.Context, step workflow.Step) (model.LLM, error) { providerID := step.Provider if providerID == "" { @@ -86,7 +87,7 @@ func (e *Engine) BuildWorkflowAgent(ctx context.Context, cwd string, def workflo if tf := GetDefaultToolFactory(); tf != nil { agentTools = tf(ToolFactoryArgs{ Cwd: cwd, - TabID: 0, + TabID: tabID, SkipPermissions: true, EventListener: e.opts.EventListener, InteractionHandler: e.opts.InteractionHandler, @@ -112,8 +113,13 @@ func (e *Engine) BuildWorkflowAgent(ctx context.Context, cwd string, def workflo } return workflow.BuildStepPrompt(step, src, nil, pc) }, + SessionService: NewFileSessionService("ask-workflow", cwd), } - return workflow.BuildWorkflowAgent(ctx, cfg) +} + +func (e *Engine) BuildWorkflowAgent(ctx context.Context, cwd string, def workflow.Def, src workflow.Source) (agent.Agent, error) { + cfg := e.BuildWorkflowAgentConfig(ctx, cwd, 0, def, src) + return workflow.CompileDefToADKWorkflow(ctx, cfg) } type engineWorkflowListener struct { @@ -159,7 +165,8 @@ func (l engineWorkflowListener) OnNote(tabID int, text string) { func (e *Engine) RunWorkflow(ctx context.Context, cwd string, tabID int, def workflow.Def, src workflow.Source) error { listener := engineWorkflowListener{tabID: tabID, listener: e.opts.EventListener} - runner := workflow.NewRunner(workflow.GlobalTracker(), e.coordinator, listener) - _, err := runner.Run(ctx, cwd, tabID, def, src) + cfg := e.BuildWorkflowAgentConfig(ctx, cwd, tabID, def, src) + runner := workflow.NewRunner(workflow.GlobalTracker(), cfg) + _, err := runner.Run(ctx, listener) return err } diff --git a/pkg/engine/plugins.go b/pkg/engine/plugins.go index 0d216ff..6b052da 100644 --- a/pkg/engine/plugins.go +++ b/pkg/engine/plugins.go @@ -17,16 +17,6 @@ func DefaultPlugins() []*plugin.Plugin { plugins = append(plugins, retryPlugin) } - // 2. Function call modifier plugin for parameter injection when configured. - // Defaults to inactive so native tools using functiontool.New with ParametersJsonSchema - // are not corrupted by the plugin's decl.Parameters initialization. - if modPlugin, err := NewFunctionCallModifierPlugin(FunctionCallModifierOptions{ - Predicate: func(toolName string) bool { - return false - }, - }); err == nil && modPlugin != nil { - plugins = append(plugins, modPlugin) - } return plugins } @@ -53,24 +43,5 @@ func NewRetryAndReflectPlugin(maxRetries int) (*plugin.Plugin, error) { } // FunctionCallModifierOptions defines configuration for the functioncallmodifier plugin. -type FunctionCallModifierOptions struct { - Predicate func(toolName string) bool - Args map[string]*genai.Schema - OverrideDescription func(originalDescription string) string -} // NewFunctionCallModifierPlugin creates an ADK functioncallmodifier plugin with safety guards. -func NewFunctionCallModifierPlugin(opts FunctionCallModifierOptions) (*plugin.Plugin, error) { - pred := opts.Predicate - if pred == nil { - pred = func(toolName string) bool { - return len(opts.Args) > 0 || opts.OverrideDescription != nil - } - } - cfg := functioncallmodifier.FunctionCallModifierConfig{ - Predicate: pred, - Args: opts.Args, - OverrideDescription: opts.OverrideDescription, - } - return functioncallmodifier.NewPlugin(cfg) -} diff --git a/pkg/engine/plugins_test.go b/pkg/engine/plugins_test.go index 2cf2e0c..c7dde64 100644 --- a/pkg/engine/plugins_test.go +++ b/pkg/engine/plugins_test.go @@ -61,40 +61,6 @@ func TestNewRetryAndReflectPlugin(t *testing.T) { }) } -func TestNewFunctionCallModifierPlugin(t *testing.T) { - t.Run("with nil predicate and args", func(t *testing.T) { - p, err := NewFunctionCallModifierPlugin(FunctionCallModifierOptions{}) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if p == nil || p.Name() != "FunctionCallModifierPlugin" { - t.Fatalf("expected valid FunctionCallModifierPlugin, got %v", p) - } - }) - - t.Run("with custom schema args and description override", func(t *testing.T) { - p, err := NewFunctionCallModifierPlugin(FunctionCallModifierOptions{ - Predicate: func(toolName string) bool { - return toolName == "read" - }, - Args: map[string]*genai.Schema{ - "description": { - Type: "STRING", - Description: "short phrase", - }, - }, - OverrideDescription: func(orig string) string { - return orig + " (modified)" - }, - }) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if p == nil || p.Name() != "FunctionCallModifierPlugin" { - t.Fatalf("expected valid FunctionCallModifierPlugin, got %v", p) - } - }) -} func TestFunctionCallModifier_ActiveParameterInjection(t *testing.T) { plugins := DefaultPlugins() diff --git a/pkg/engine/prompt.go b/pkg/engine/prompt.go index db87299..3f6ed10 100644 --- a/pkg/engine/prompt.go +++ b/pkg/engine/prompt.go @@ -496,22 +496,6 @@ func BuildSystemPrompt(opts PromptOptions) string { } } - if memory.IsOpen() { - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - if mem := memory.SystemBlock(ctx, cwd); mem != "" { - b.WriteString("\n\n\n") - b.WriteString(mem) - b.WriteString("\n") - } - cancel() - } - - if !opts.DisableSkillsPrompt { - if block := SkillsPromptBlock(DiscoverSkills(cwd)); block != "" { - b.WriteString("\n\n") - b.WriteString(block) - } - } if block := SubagentsPromptBlock(DiscoverSubagents(cwd)); block != "" { b.WriteString("\n\n") b.WriteString(block) diff --git a/pkg/engine/types.go b/pkg/engine/types.go index 69b593a..f1ff8cc 100644 --- a/pkg/engine/types.go +++ b/pkg/engine/types.go @@ -3,11 +3,14 @@ package engine import ( "context" "encoding/json" + "errors" "fmt" "strings" "time" + pkgmemory "github.com/Cidan/ask/pkg/memory" "github.com/google/jsonschema-go/jsonschema" + "google.golang.org/adk/v2/artifact" "google.golang.org/adk/v2/agent" "google.golang.org/adk/v2/memory" "google.golang.org/adk/v2/model" @@ -16,7 +19,6 @@ import ( "google.golang.org/adk/v2/tool/functiontool" "google.golang.org/adk/v2/tool/toolconfirmation" "google.golang.org/genai" - pkgmemory "github.com/Cidan/ask/pkg/memory" ) // ToolResponse represents the result of executing a tool. @@ -57,7 +59,9 @@ func ExtractToolInfo(t tool.Tool) ToolInfo { Name: t.Name(), Description: t.Description(), } - if declProvider, ok := t.(interface{ Declaration() *genai.FunctionDeclaration }); ok { + if declProvider, ok := t.(interface { + Declaration() *genai.FunctionDeclaration + }); ok { if decl := declProvider.Declaration(); decl != nil { if decl.ParametersJsonSchema != nil { if raw, err := json.Marshal(decl.ParametersJsonSchema); err == nil { @@ -92,30 +96,34 @@ func ExtractToolInfo(t tool.Tool) ToolInfo { type standaloneAgentContext struct { context.Context + artifactService artifact.Service } -func (s *standaloneAgentContext) UserContent() *genai.Content { return nil } -func (s *standaloneAgentContext) InvocationID() string { return "" } -func (s *standaloneAgentContext) AgentName() string { return "" } -func (s *standaloneAgentContext) ReadonlyState() session.ReadonlyState { return nil } -func (s *standaloneAgentContext) UserID() string { return "" } -func (s *standaloneAgentContext) AppName() string { return "ask" } -func (s *standaloneAgentContext) SessionID() string { return "" } -func (s *standaloneAgentContext) Branch() string { return "" } -func (s *standaloneAgentContext) Agent() agent.Agent { return nil } -func (s *standaloneAgentContext) Artifacts() agent.Artifacts { return nil } -func (s *standaloneAgentContext) Memory() agent.Memory { return nil } -func (s *standaloneAgentContext) Session() session.Session { return nil } -func (s *standaloneAgentContext) IsolationScope() string { return "" } -func (s *standaloneAgentContext) RunConfig() *agent.RunConfig { return nil } -func (s *standaloneAgentContext) EndInvocation() {} -func (s *standaloneAgentContext) Ended() bool { return false } -func (s *standaloneAgentContext) ResumedInput(interruptID string) (any, bool) { return nil, false } -func (s *standaloneAgentContext) WithContext(ctx context.Context) agent.InvocationContext { return &standaloneAgentContext{Context: ctx} } -func (s *standaloneAgentContext) WithICDelta(d *agent.InvocationContextDelta) agent.InvocationContext { return s } -func (s *standaloneAgentContext) State() session.State { return nil } -func (s *standaloneAgentContext) FunctionCallID() string { return "" } -func (s *standaloneAgentContext) Actions() *session.EventActions { return &session.EventActions{} } +func (s *standaloneAgentContext) UserContent() *genai.Content { return nil } +func (s *standaloneAgentContext) InvocationID() string { return "" } +func (s *standaloneAgentContext) AgentName() string { return "" } +func (s *standaloneAgentContext) ReadonlyState() session.ReadonlyState { return nil } +func (s *standaloneAgentContext) UserID() string { return "" } +func (s *standaloneAgentContext) AppName() string { return "ask" } +func (s *standaloneAgentContext) SessionID() string { return "" } +func (s *standaloneAgentContext) Branch() string { return "" } +func (s *standaloneAgentContext) Agent() agent.Agent { return nil } +func (s *standaloneAgentContext) Memory() agent.Memory { return nil } +func (s *standaloneAgentContext) Session() session.Session { return nil } +func (s *standaloneAgentContext) IsolationScope() string { return "" } +func (s *standaloneAgentContext) RunConfig() *agent.RunConfig { return nil } +func (s *standaloneAgentContext) EndInvocation() {} +func (s *standaloneAgentContext) Ended() bool { return false } +func (s *standaloneAgentContext) ResumedInput(interruptID string) (any, bool) { return nil, false } +func (s *standaloneAgentContext) WithContext(ctx context.Context) agent.InvocationContext { + return &standaloneAgentContext{Context: ctx} +} +func (s *standaloneAgentContext) WithICDelta(d *agent.InvocationContextDelta) agent.InvocationContext { + return s +} +func (s *standaloneAgentContext) State() session.State { return nil } +func (s *standaloneAgentContext) FunctionCallID() string { return "" } +func (s *standaloneAgentContext) Actions() *session.EventActions { return &session.EventActions{} } func (s *standaloneAgentContext) SearchMemory(ctx context.Context, query string) (*memory.SearchResponse, error) { if memSvc := pkgmemory.Default(); memSvc != nil && memSvc.IsOpen() { return memSvc.SearchMemory(ctx, &memory.SearchRequest{Query: query}) @@ -128,10 +136,12 @@ func (s *standaloneAgentContext) ToolConfirmation() *toolconfirmation.ToolConfir func (s *standaloneAgentContext) RequestConfirmation(hint string, payload any) error { return nil } -func (s *standaloneAgentContext) Path() string { return "" } -func (s *standaloneAgentContext) RunID() string { return "" } -func (s *standaloneAgentContext) SubScheduler() agent.DynamicSubScheduler { return nil } -func (s *standaloneAgentContext) WithAgentContext(ctx context.Context) agent.Context { return &standaloneAgentContext{Context: ctx} } +func (s *standaloneAgentContext) Path() string { return "" } +func (s *standaloneAgentContext) RunID() string { return "" } +func (s *standaloneAgentContext) SubScheduler() agent.DynamicSubScheduler { return nil } +func (s *standaloneAgentContext) WithAgentContext(ctx context.Context) agent.Context { + return &standaloneAgentContext{Context: ctx} +} func (s *standaloneAgentContext) WithAgentTimeout(timeout time.Duration) (agent.Context, context.CancelFunc) { ctx, cancel := context.WithTimeout(s.Context, timeout) return &standaloneAgentContext{Context: ctx}, cancel @@ -140,8 +150,8 @@ func (s *standaloneAgentContext) WithAgentCancel() (agent.Context, context.Cance ctx, cancel := context.WithCancel(s.Context) return &standaloneAgentContext{Context: ctx}, cancel } -func (s *standaloneAgentContext) OutputForAncestors() []string { return nil } -func (s *standaloneAgentContext) WithDelta(d *agent.CommonContextDelta) agent.Context { return s } +func (s *standaloneAgentContext) OutputForAncestors() []string { return nil } +func (s *standaloneAgentContext) WithDelta(d *agent.CommonContextDelta) agent.Context { return s } // NewStandaloneAgentContext wraps a context.Context with a compliant agent.Context implementation. func NewStandaloneAgentContext(ctx context.Context) agent.Context { @@ -569,3 +579,70 @@ func (m *Message) UnmarshalJSON(data []byte) error { } return nil } + +type standaloneArtifacts struct { + service artifact.Service + sessionID string + appID string + userID string +} + +func (a *standaloneArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { + if a.service == nil { + return nil, errors.New("no artifact service configured") + } + return a.service.Save(ctx, &artifact.SaveRequest{ + AppName: a.appID, + UserID: a.userID, + SessionID: a.sessionID, + FileName: name, + Data: data, + }) +} + +func (a *standaloneArtifacts) Load(ctx context.Context, name string) (*artifact.LoadResponse, error) { + if a.service == nil { + return nil, errors.New("no artifact service configured") + } + return a.service.Load(ctx, &artifact.LoadRequest{ + AppName: a.appID, + UserID: a.userID, + SessionID: a.sessionID, + FileName: name, + }) +} + +func (a *standaloneArtifacts) LoadVersion(ctx context.Context, name string, version int) (*artifact.LoadResponse, error) { + if a.service == nil { + return nil, errors.New("no artifact service configured") + } + return a.service.Load(ctx, &artifact.LoadRequest{ + AppName: a.appID, + UserID: a.userID, + SessionID: a.sessionID, + FileName: name, + }) +} + +func (a *standaloneArtifacts) List(ctx context.Context) (*artifact.ListResponse, error) { + if a.service == nil { + return nil, errors.New("no artifact service configured") + } + return a.service.List(ctx, &artifact.ListRequest{ + AppName: a.appID, + UserID: a.userID, + SessionID: a.sessionID, + }) +} + +func (s *standaloneAgentContext) Artifacts() agent.Artifacts { + if s.artifacts == nil && s.artifactService != nil { + s.artifacts = &standaloneArtifacts{ + service: s.artifactService, + sessionID: s.SessionID(), + appID: s.AppName(), + userID: s.UserID(), + } + } + return s.artifacts +} diff --git a/pkg/tools/core.go b/pkg/tools/core.go index d22a78a..c283ef2 100644 --- a/pkg/tools/core.go +++ b/pkg/tools/core.go @@ -104,3 +104,6 @@ func IsCoreTool(name string) bool { return false } } +func LoadArtifactsTool() Tool { + return loadartifactstool.New() +} diff --git a/pkg/tools/memory.go b/pkg/tools/memory.go index f645fdc..ee17869 100644 --- a/pkg/tools/memory.go +++ b/pkg/tools/memory.go @@ -2,15 +2,11 @@ package tools import ( "context" - "encoding/json" "strings" - "github.com/Cidan/ask/pkg/engine" "github.com/Cidan/ask/pkg/memory" - "google.golang.org/adk/v2/agent" - "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/tool/loadmemorytool" "google.golang.org/adk/v2/tool/preloadmemorytool" - "google.golang.org/genai" ) type loadMemoryParams struct { @@ -19,38 +15,7 @@ type loadMemoryParams struct { // LoadMemoryTool returns the native tool for querying long-term memory with JSON Schema compatibility. func LoadMemoryTool() Tool { - return NewTool( - "load_memory", - "Loads the memory for the current user.", - func(ctx context.Context, p loadMemoryParams) (ToolResponse, error) { - query := strings.TrimSpace(p.Query) - if query == "" { - return NewTextErrorResponse("query cannot be empty"), nil - } - actx := engine.NewStandaloneAgentContext(ctx) - res, err := actx.SearchMemory(ctx, query) - if err != nil { - return NewTextErrorResponse("failed to search memory: " + err.Error()), nil - } - if res == nil || len(res.Memories) == 0 { - return NewTextResponse("no memories found"), nil - } - var lines []string - for _, m := range res.Memories { - if m.Content != nil { - for _, part := range m.Content.Parts { - if part.Text != "" { - lines = append(lines, part.Text) - } - } - } - } - if len(lines) == 0 { - return NewTextResponse("no memories found"), nil - } - return NewTextResponse(strings.Join(lines, "\n")), nil - }, - ) + return loadmemorytool.New() } // PreloadMemoryTool returns the native ADK preload_memory tool for automatic turn recall. @@ -94,85 +59,7 @@ func MemoryIndexTool(cwd string, requestApproval func(ctx context.Context, name // MemoryAwareTool decorates a file tool (read / edit / write) so the tool output carries // relevant memory recall context for the touched file path. -type MemoryAwareTool struct { - Inner Tool - Cwd string -} - -func (m *MemoryAwareTool) Name() string { return m.Inner.Name() } -func (m *MemoryAwareTool) Description() string { return m.Inner.Description() } -func (m *MemoryAwareTool) IsLongRunning() bool { return m.Inner.IsLongRunning() } -func (m *MemoryAwareTool) Info() ToolInfo { return ExtractToolInfo(m.Inner) } -func (m *MemoryAwareTool) Declaration() *genai.FunctionDeclaration { - if dp, ok := m.Inner.(interface{ Declaration() *genai.FunctionDeclaration }); ok { - return dp.Declaration() - } - return nil -} - -func (m *MemoryAwareTool) ProcessRequest(ctx agent.Context, req *model.LLMRequest) error { - if rp, ok := m.Inner.(interface { - ProcessRequest(ctx agent.Context, req *model.LLMRequest) error - }); ok { - return rp.ProcessRequest(ctx, req) - } - return nil -} // WrapFileToolsWithMemory decorates read, edit, and write tools with memory recall. -func WrapFileToolsWithMemory(tools []Tool, cwd string) []Tool { - out := make([]Tool, len(tools)) - for i, t := range tools { - switch t.Name() { - case "read", "edit", "write": - out[i] = &MemoryAwareTool{Inner: t, Cwd: cwd} - default: - out[i] = t - } - } - return out -} // Run executes the underlying tool and appends file-specific memory recall when available. -func (m *MemoryAwareTool) Run(ctx agent.Context, args any) (map[string]any, error) { - resp, err := RunADKTool(ctx, m.Inner, args) - if err != nil || !memory.IsOpen() { - return resp, err - } - if isErr, _ := resp["is_error"].(bool); isErr { - return resp, err - } - - argsMap, _ := args.(map[string]any) - if argsMap == nil { - if raw, err := json.Marshal(args); err == nil { - _ = json.Unmarshal(raw, &argsMap) - } - } - - path, _ := argsMap["file_path"].(string) - path = strings.TrimSpace(path) - if path == "" { - return resp, err - } - recallCtx, cancel := context.WithTimeout(ctx, memory.DefaultHookTimeout) - defer cancel() - hits, rerr := memory.Recall(recallCtx, m.Cwd, path, memory.DefaultRecallK) - if rerr != nil { - return resp, err - } - block := memory.FormatRecallContext(hits, "Memory for "+path) - if block == "" { - return resp, err - } - - if resp == nil { - resp = make(map[string]any) - } - if s, ok := resp["result"].(string); ok && s != "" { - resp["result"] = s + "\n\n" + block - } else { - resp["result"] = block - } - return resp, err -} diff --git a/pkg/workflow/agent_factory.go b/pkg/workflow/agent_factory.go new file mode 100644 index 0000000..354f388 --- /dev/null +++ b/pkg/workflow/agent_factory.go @@ -0,0 +1,15 @@ +package workflow + +import ( + "context" + + "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/tool" +) + +type AgentFactory interface { + ModelBuilder(ctx context.Context, step Step) (model.LLM, error) + ToolsBuilder(ctx context.Context, step Step, isLoop bool) ([]tool.Tool, error) + ToolsetsBuilder(ctx context.Context, step Step, isLoop bool) ([]tool.Toolset, error) + InstructionBuilder(step Step, isStart bool, isFinal bool, loopCtx *LoopPromptCtx, notesDir, prevNotesDir string) string +} diff --git a/pkg/workflow/graph.go b/pkg/workflow/graph.go index 0083d72..90b8261 100644 --- a/pkg/workflow/graph.go +++ b/pkg/workflow/graph.go @@ -7,15 +7,16 @@ import ( "google.golang.org/adk/v2/agent" "google.golang.org/adk/v2/agent/llmagent" + "google.golang.org/adk/v2/agent/workflowagent" "google.golang.org/adk/v2/agent/workflowagents/loopagent" "google.golang.org/adk/v2/tool" "google.golang.org/adk/v2/tool/exitlooptool" adkworkflow "google.golang.org/adk/v2/workflow" ) -// CompileDefToADKWorkflow converts a workflow.Def into an ADK 2.0 directed acyclic graph (*adkworkflow.Workflow). +// CompileDefToADKWorkflow converts a workflow.Def into an ADK 2.0 directed acyclic graph (agent.Agent). // It constructs agent nodes for each top-level step and connects them using standard workflow edges and routes. -func CompileDefToADKWorkflow(ctx context.Context, cfg WorkflowAgentConfig) (*adkworkflow.Workflow, error) { +func CompileDefToADKWorkflow(ctx context.Context, cfg WorkflowAgentConfig) (agent.Agent, error) { if err := cfg.Def.Validate(); err != nil { return nil, err } @@ -24,6 +25,7 @@ func CompileDefToADKWorkflow(ctx context.Context, cfg WorkflowAgentConfig) (*adk } var nodes []adkworkflow.Node + var subAgents []agent.Agent var prevNotesDir string for i, top := range cfg.Def.Steps { @@ -125,6 +127,7 @@ func CompileDefToADKWorkflow(ctx context.Context, cfg WorkflowAgentConfig) (*adk return nil, fmt.Errorf("failed to create agent node for loop %q: %w", top.Name, err) } nodes = append(nodes, loopNode) + subAgents = append(subAgents, loopAg) continue } @@ -182,6 +185,7 @@ func CompileDefToADKWorkflow(ctx context.Context, cfg WorkflowAgentConfig) (*adk return nil, fmt.Errorf("failed to create agent node for step %q: %w", top.Name, err) } nodes = append(nodes, stepNode) + subAgents = append(subAgents, stepAg) prevNotesDir = notesDir } @@ -190,5 +194,10 @@ func CompileDefToADKWorkflow(ctx context.Context, cfg WorkflowAgentConfig) (*adk } edges := append([]adkworkflow.Edge{{From: adkworkflow.Start, To: nodes[0]}}, adkworkflow.Chain(nodes...)...) - return adkworkflow.New(cfg.Def.Name, edges) + return workflowagent.New(workflowagent.Config{ + Name: cfg.Def.Name, + Description: cfg.Def.Description, + SubAgents: subAgents, + Edges: edges, + }) } diff --git a/pkg/workflow/graph_test.go b/pkg/workflow/graph_test.go index b916f8b..4d06bc1 100644 --- a/pkg/workflow/graph_test.go +++ b/pkg/workflow/graph_test.go @@ -134,67 +134,7 @@ func TestCompileDefToADKWorkflow_ValidationErrors(t *testing.T) { } } -type testGraphListener struct { - NoopRunnerListener - started bool - stepsDone int - done bool -} - -func (l *testGraphListener) OnWorkflowStarted(tabID int, def Def, src Source) { - l.started = true -} -func (l *testGraphListener) OnWorkflowStepDone(tabID int, stepIdx int, summary string) { - l.stepsDone++ -} -func (l *testGraphListener) OnWorkflowDone(tabID int, desc string, artifacts []string) { - l.done = true -} - -func TestWorkflowRunner_ADKGraphExecution(t *testing.T) { - def := Def{ - Name: "adk-graph-run", - Description: "executing workflow via adk graph", - Steps: []Step{ - {Name: "step-1", Prompt: "do step 1", Provider: "vertex", Model: "gemini"}, - {Name: "step-2", Prompt: "do step 2", Provider: "vertex", Model: "gemini"}, - }, - } - listener := &testGraphListener{} - runner := NewRunner(NewTracker(), nil, listener) - cfg := WorkflowAgentConfig{ - Def: def, - Cwd: t.TempDir(), - TabID: 42, - ModelBuilder: func(ctx context.Context, step Step) (model.LLM, error) { - return &fakeModel{}, nil - }, - ToolsBuilder: func(ctx context.Context, step Step, isLoop bool) ([]tool.Tool, error) { - return nil, nil - }, - } - - state, err := runner.RunGraph(context.Background(), cfg) - if err != nil { - t.Fatalf("unexpected error running adk graph workflow: %v", err) - } - if state == nil || !state.Done { - t.Fatalf("expected completed run state, got %+v", state) - } - if state.StepIdx != 2 { - t.Errorf("expected StepIdx 2, got %d", state.StepIdx) - } - if !listener.started { - t.Error("expected listener OnWorkflowStarted to be called") - } - if listener.stepsDone != 2 { - t.Errorf("expected 2 step done calls, got %d", listener.stepsDone) - } - if !listener.done { - t.Error("expected listener OnWorkflowDone to be called") - } -} diff --git a/pkg/workflow/runner.go b/pkg/workflow/runner.go index dbf7abd..b803698 100644 --- a/pkg/workflow/runner.go +++ b/pkg/workflow/runner.go @@ -2,21 +2,17 @@ package workflow import ( "context" - "errors" "fmt" "path/filepath" "strings" "time" "google.golang.org/adk/v2/agent" - "google.golang.org/adk/v2/agent/llmagent" - "google.golang.org/adk/v2/agent/workflowagents/loopagent" - "google.golang.org/adk/v2/agent/workflowagents/sequentialagent" + "google.golang.org/adk/v2/artifact" "google.golang.org/adk/v2/model" "google.golang.org/adk/v2/runner" "google.golang.org/adk/v2/session" "google.golang.org/adk/v2/tool" - "google.golang.org/adk/v2/tool/exitlooptool" "google.golang.org/genai" ) @@ -95,9 +91,6 @@ type StepResult struct { } // StepExecutor executes a single step turn against an underlying agent engine/provider. -type StepExecutor interface { - ExecuteStep(ctx context.Context, cwd string, tabID int, step Step, prompt string, isFinal bool) (StepResult, error) -} // RunnerListener receives progress notifications during workflow execution. type RunnerListener interface { @@ -121,6 +114,8 @@ func (NoopRunnerListener) OnNote(int, string) // WorkflowAgentConfig configures the construction of an ADK workflow agent hierarchy. type WorkflowAgentConfig struct { + SessionService session.Service + ArtifactService artifact.Service Def Def Source Source Cwd string @@ -131,243 +126,73 @@ type WorkflowAgentConfig struct { InstructionBuilder func(step Step, isStart bool, isFinal bool, loopCtx *LoopPromptCtx, notesDir, prevNotesDir string) string } -// BuildWorkflowAgent constructs an ADK agent hierarchy conforming to the workflow definition. -// Top-level linear steps are chained using sequentialagent, while kind: "loop" steps are -// encapsulated in loopagent containers with exitlooptool attached to their sub-agents. -func BuildWorkflowAgent(ctx context.Context, cfg WorkflowAgentConfig) (agent.Agent, error) { - if err := cfg.Def.Validate(); err != nil { - return nil, err - } - if cfg.ModelBuilder == nil { - return nil, errors.New("model builder is required") - } - - var topAgents []agent.Agent - var prevNotesDir string - - for i, top := range cfg.Def.Steps { - isFinalStep := i == len(cfg.Def.Steps)-1 - - if top.Kind == "loop" { - if len(top.Steps) == 0 { - continue - } - var innerAgents []agent.Agent - for innerIdx, innerStep := range top.Steps { - isLoopStart := i == 0 && innerIdx == 0 - var notesDir string - if isLoopStart { - notesDir = StartPlanDir(cfg.Cwd) - } else { - notesDir = StepNotesDir(cfg.Cwd, innerStep.Name, top.Name, 1) - } - - llm, err := cfg.ModelBuilder(ctx, innerStep) - if err != nil { - return nil, fmt.Errorf("failed to build model for step %q: %w", innerStep.Name, err) - } - - var tools []tool.Tool - if cfg.ToolsBuilder != nil { - builtTools, err := cfg.ToolsBuilder(ctx, innerStep, true) - if err != nil { - return nil, fmt.Errorf("failed to build tools for step %q: %w", innerStep.Name, err) - } - tools = append(tools, builtTools...) - } - - // Attach ADK's native exitlooptool for clean early break out of loop containers - exitTool, err := exitlooptool.New() - if err != nil { - return nil, fmt.Errorf("failed to create exitloop tool: %w", err) - } - hasExitTool := false - for _, t := range tools { - if t != nil && t.Name() == exitTool.Name() { - hasExitTool = true - break - } - } - if !hasExitTool { - tools = append(tools, exitTool) - } - - var toolsets []tool.Toolset - if cfg.ToolsetsBuilder != nil { - ts, err := cfg.ToolsetsBuilder(ctx, innerStep, true) - if err != nil { - return nil, fmt.Errorf("failed to build toolsets for step %q: %w", innerStep.Name, err) - } - toolsets = ts - } - - instruction := innerStep.Prompt - if cfg.InstructionBuilder != nil { - loopCtx := &LoopPromptCtx{ - Name: top.Name, - Iteration: 1, - MaxIterations: cfg.Def.EffectiveMaxIterations(top), - ExitCondition: top.ExitCondition, - IsTail: innerIdx == len(top.Steps)-1, - } - instruction = cfg.InstructionBuilder(innerStep, isLoopStart, isFinalStep, loopCtx, notesDir, prevNotesDir) - } - - innerAg, err := llmagent.New(llmagent.Config{ - Name: innerStep.Name, - Description: innerStep.Prompt, - Model: llm, - Instruction: instruction, - Tools: tools, - Toolsets: toolsets, - }) - if err != nil { - return nil, fmt.Errorf("failed to create inner step agent %q: %w", innerStep.Name, err) - } - innerAgents = append(innerAgents, innerAg) - prevNotesDir = notesDir - } - - loopAg, err := loopagent.New(loopagent.Config{ - AgentConfig: agent.Config{ - Name: top.Name, - Description: top.ExitCondition, - SubAgents: innerAgents, - }, - MaxIterations: uint(cfg.Def.EffectiveMaxIterations(top)), - }) - if err != nil { - return nil, fmt.Errorf("failed to create loop agent %q: %w", top.Name, err) - } - topAgents = append(topAgents, loopAg) - continue - } - - // Linear step - isStart := i == 0 - var notesDir string - if isStart { - notesDir = StartPlanDir(cfg.Cwd) - } else { - notesDir = StepNotesDir(cfg.Cwd, top.Name, "", 0) - } - - llm, err := cfg.ModelBuilder(ctx, top) - if err != nil { - return nil, fmt.Errorf("failed to build model for step %q: %w", top.Name, err) - } - - var tools []tool.Tool - if cfg.ToolsBuilder != nil { - builtTools, err := cfg.ToolsBuilder(ctx, top, false) - if err != nil { - return nil, fmt.Errorf("failed to build tools for step %q: %w", top.Name, err) - } - tools = append(tools, builtTools...) - } - - var toolsets []tool.Toolset - if cfg.ToolsetsBuilder != nil { - ts, err := cfg.ToolsetsBuilder(ctx, top, false) - if err != nil { - return nil, fmt.Errorf("failed to build toolsets for step %q: %w", top.Name, err) - } - toolsets = ts - } - - instruction := top.Prompt - if cfg.InstructionBuilder != nil { - instruction = cfg.InstructionBuilder(top, isStart, isFinalStep, nil, notesDir, prevNotesDir) - } +// Runner executes multi-step workflow pipelines. - stepAg, err := llmagent.New(llmagent.Config{ - Name: top.Name, - Description: top.Prompt, - Model: llm, - Instruction: instruction, - Tools: tools, - Toolsets: toolsets, - }) - if err != nil { - return nil, fmt.Errorf("failed to create step agent %q: %w", top.Name, err) - } - topAgents = append(topAgents, stepAg) - prevNotesDir = notesDir - } +// NewRunner creates a new workflow Runner. - return sequentialagent.New(sequentialagent.Config{ - AgentConfig: agent.Config{ - Name: cfg.Def.Name, - Description: cfg.Def.Description, - SubAgents: topAgents, - }, - }) -} +// Run executes the workflow def synchronously to completion or until context cancellation. -// Runner executes multi-step workflow pipelines. type Runner struct { - tracker *Tracker - executor StepExecutor - listener RunnerListener + cfg WorkflowAgentConfig + tracker *Tracker } -// NewRunner creates a new workflow Runner. -func NewRunner(tracker *Tracker, executor StepExecutor, listener RunnerListener) *Runner { +func NewRunner(tracker *Tracker, cfg WorkflowAgentConfig) *Runner { if tracker == nil { tracker = GlobalTracker() } - if listener == nil { - listener = NoopRunnerListener{} - } return &Runner{ - tracker: tracker, - executor: executor, - listener: listener, + cfg: cfg, + tracker: tracker, } } -// RunGraph executes a compiled ADK workflow graph, broadcasting lifecycle events to the listener. -func (r *Runner) RunGraph(ctx context.Context, cfg WorkflowAgentConfig) (*RunState, error) { - if err := cfg.Def.Validate(); err != nil { - r.listener.OnWorkflowFailed(cfg.TabID, err.Error()) +func (r *Runner) Run(ctx context.Context, listener RunnerListener) (*RunState, error) { + if err := r.cfg.Def.Validate(); err != nil { + listener.OnWorkflowFailed(r.cfg.TabID, err.Error()) return nil, err } - wfAgent, err := BuildWorkflowAgent(ctx, cfg) + wfAgent, err := CompileDefToADKWorkflow(ctx, r.cfg) if err != nil { - r.listener.OnWorkflowFailed(cfg.TabID, err.Error()) + listener.OnWorkflowFailed(r.cfg.TabID, err.Error()) return nil, err } - r.listener.OnWorkflowStarted(cfg.TabID, cfg.Def, cfg.Source) + listener.OnWorkflowStarted(r.cfg.TabID, r.cfg.Def, r.cfg.Source) if r.tracker != nil { - r.tracker.MarkWorking(cfg.Cwd, cfg.Source.Key(), cfg.Def.Name, cfg.TabID) + r.tracker.MarkWorking(r.cfg.Cwd, r.cfg.Source.Key(), r.cfg.Def.Name, r.cfg.TabID) } runState := &RunState{ - Workflow: cfg.Def, - Source: cfg.Source, + Workflow: r.cfg.Def, + Source: r.cfg.Source, StartedAt: time.Now().UTC(), StepIdx: 0, } - sessSvc := session.InMemoryService() + sessSvc := r.cfg.SessionService + if sessSvc == nil { + sessSvc = session.InMemoryService() + } + adkRunner, err := runner.New(runner.Config{ AppName: "ask-workflow", Agent: wfAgent, SessionService: sessSvc, + ArtifactService: r.cfg.ArtifactService, AutoCreateSession: true, }) if err != nil { - r.listener.OnWorkflowFailed(cfg.TabID, err.Error()) + listener.OnWorkflowFailed(r.cfg.TabID, err.Error()) if r.tracker != nil { - r.tracker.MarkFinal(cfg.Cwd, cfg.Source.Key(), cfg.Def.Name, StatusFailed, 0) + r.tracker.MarkFinal(r.cfg.Cwd, r.cfg.Source.Key(), r.cfg.Def.Name, StatusFailed, 0) } return runState, err } - userMsg := genai.NewContentFromText(cfg.Source.Display(), genai.RoleUser) - sessionID := "wf-" + cfg.Source.Key() + userMsg := genai.NewContentFromText(r.cfg.Source.Display(), genai.RoleUser) + sessionID := "wf-" + r.cfg.Source.Key() startedSteps := make(map[int]bool) doneSteps := make(map[int]bool) @@ -375,9 +200,9 @@ func (r *Runner) RunGraph(ctx context.Context, cfg WorkflowAgentConfig) (*RunSta for event, err := range adkRunner.Run(ctx, "user", sessionID, userMsg, agent.RunConfig{}) { if err != nil { - r.listener.OnWorkflowFailed(cfg.TabID, err.Error()) + listener.OnWorkflowFailed(r.cfg.TabID, err.Error()) if r.tracker != nil { - r.tracker.MarkFinal(cfg.Cwd, cfg.Source.Key(), cfg.Def.Name, StatusFailed, lastStepIdx) + r.tracker.MarkFinal(r.cfg.Cwd, r.cfg.Source.Key(), r.cfg.Def.Name, StatusFailed, lastStepIdx) } return runState, err } @@ -386,16 +211,16 @@ func (r *Runner) RunGraph(ctx context.Context, cfg WorkflowAgentConfig) (*RunSta } if event.Author != "" && event.Author != "user" && event.Author != "ask_coder" { - for i, s := range cfg.Def.Steps { + for i, s := range r.cfg.Def.Steps { if s.Name == event.Author { if lastStepIdx >= 0 && lastStepIdx != i && !doneSteps[lastStepIdx] { doneSteps[lastStepIdx] = true - r.listener.OnWorkflowStepDone(cfg.TabID, lastStepIdx, fmt.Sprintf("completed step %s", cfg.Def.Steps[lastStepIdx].Name)) + listener.OnWorkflowStepDone(r.cfg.TabID, lastStepIdx, fmt.Sprintf("completed step %s", r.cfg.Def.Steps[lastStepIdx].Name)) } if !startedSteps[i] { startedSteps[i] = true lastStepIdx = i - r.listener.OnWorkflowStepStarted(cfg.TabID, i, s.Name, s.Provider, s.Model) + listener.OnWorkflowStepStarted(r.cfg.TabID, i, s.Name, s.Provider, s.Model) } break } @@ -405,317 +230,32 @@ func (r *Runner) RunGraph(ctx context.Context, cfg WorkflowAgentConfig) (*RunSta if event.LLMResponse.Content != nil { for _, part := range event.LLMResponse.Content.Parts { if part.Text != "" { - r.listener.OnNote(cfg.TabID, part.Text) + listener.OnNote(r.cfg.TabID, part.Text) } } } } - for i := 0; i < len(cfg.Def.Steps); i++ { + for i := 0; i < len(r.cfg.Def.Steps); i++ { if !startedSteps[i] { startedSteps[i] = true - r.listener.OnWorkflowStepStarted(cfg.TabID, i, cfg.Def.Steps[i].Name, cfg.Def.Steps[i].Provider, cfg.Def.Steps[i].Model) + listener.OnWorkflowStepStarted(r.cfg.TabID, i, r.cfg.Def.Steps[i].Name, r.cfg.Def.Steps[i].Provider, r.cfg.Def.Steps[i].Model) } if !doneSteps[i] { doneSteps[i] = true - r.listener.OnWorkflowStepDone(cfg.TabID, i, fmt.Sprintf("completed step %s", cfg.Def.Steps[i].Name)) + listener.OnWorkflowStepDone(r.cfg.TabID, i, fmt.Sprintf("completed step %s", r.cfg.Def.Steps[i].Name)) } } runState.Done = true - runState.StepIdx = len(cfg.Def.Steps) - runState.FinishData = &FinishData{ - Description: fmt.Sprintf("Workflow %s completed via ADK workflow runner", cfg.Def.Name), - } - - r.listener.OnWorkflowDone(cfg.TabID, runState.FinishData.Description, runState.FinishData.Artifacts) if r.tracker != nil { - r.tracker.MarkFinal(cfg.Cwd, cfg.Source.Key(), cfg.Def.Name, StatusDone, len(cfg.Def.Steps)) + r.tracker.MarkFinal(r.cfg.Cwd, r.cfg.Source.Key(), r.cfg.Def.Name, StatusDone, len(r.cfg.Def.Steps)-1) } - return runState, nil -} - -// Run executes the workflow def synchronously to completion or until context cancellation. -func (r *Runner) Run(ctx context.Context, cwd string, tabID int, def Def, src Source) (*RunState, error) { - if err := def.Validate(); err != nil { - r.listener.OnWorkflowFailed(tabID, err.Error()) - return nil, err - } - - r.listener.OnWorkflowStarted(tabID, def, src) - r.tracker.MarkWorking(cwd, src.Key(), def.Name, tabID) - - runState := &RunState{ - Workflow: def, - Source: src, - StartedAt: time.Now().UTC(), - StepIdx: 0, - } - - var stepLog []string - var loopFrame *LoopRunFrame - var prevNotesDir string - var currentNotesDir string - var remind RemindKind - var remindDetail string - var linearRetry int - var linearText string - var stepErrorRetry int - - for { - select { - case <-ctx.Done(): - r.listener.OnWorkflowFailed(tabID, "cancelled by user") - r.tracker.MarkFinal(cwd, src.Key(), def.Name, StatusFailed, runState.StepIdx) - return runState, ctx.Err() - default: - } - - if loopFrame == nil && runState.StepIdx >= len(def.Steps) { - break - } - - top := def.Steps[runState.StepIdx] - if top.Kind == "loop" && loopFrame == nil { - if len(top.Steps) == 0 { - runState.StepIdx++ - continue - } - loopFrame = &LoopRunFrame{InnerIdx: 0, Iteration: 1} - r.listener.OnNote(tabID, LoopNoteLine(top.Name, "started", fmt.Sprintf("max %d iteration(s)", def.EffectiveMaxIterations(top)))) - } - - step := top - if loopFrame != nil { - step = top.Steps[loopFrame.InnerIdx] - } - - r.listener.OnWorkflowStepStarted(tabID, runState.StepIdx, step.Name, step.Provider, step.Model) - - isStartStep := runState.StepIdx == 0 && loopFrame == nil - isLoopStartStep := runState.StepIdx == 0 && loopFrame != nil && loopFrame.Iteration == 1 && loopFrame.InnerIdx == 0 - - var notesDir string - switch { - case isStartStep, isLoopStartStep: - notesDir = StartPlanDir(cwd) - case loopFrame != nil: - notesDir = StepNotesDir(cwd, step.Name, top.Name, loopFrame.Iteration) - default: - notesDir = StepNotesDir(cwd, step.Name, "", 0) - } - currentNotesDir = notesDir - - var prevOutputs []string - if loopFrame == nil { - if linearRetry > 0 && linearText != "" { - prevOutputs = append(append([]string(nil), stepLog...), linearText) - } else { - prevOutputs = stepLog - } - } else { - prevOutputs = append([]string(nil), stepLog...) - if loopFrame.InnerIdx == 0 { - if loopFrame.PrevTail != "" { - prevOutputs = append(prevOutputs, loopFrame.PrevTail) - } - } else { - prevOutputs = append(prevOutputs, loopFrame.IterationLog...) - } - if loopFrame.Retry > 0 && loopFrame.RetryText != "" { - prevOutputs = append(prevOutputs, loopFrame.RetryText) - } - } - - pc := &StepPromptCtx{ - Remind: remind, - RemindDetail: remindDetail, - NotesDir: notesDir, - PrevNotesDir: prevNotesDir, - IsStartStep: isStartStep || isLoopStartStep, - IsWorkflowFinalStep: runState.StepIdx == len(def.Steps)-1, - } - if loopFrame != nil { - pc.Loop = &LoopPromptCtx{ - Name: top.Name, - Iteration: loopFrame.Iteration, - MaxIterations: def.EffectiveMaxIterations(top), - ExitCondition: top.ExitCondition, - IsTail: loopFrame.InnerIdx == len(top.Steps)-1, - } - } - - var dirErr error - if pc.IsStartStep { - dirErr = EnsureStartPlanExists(cwd) - } else { - dirErr = EnsureStepNotesDir(notesDir) - } - if dirErr != nil { - remind = RemindFixPlanDir - remindDetail = dirErr.Error() - pc.Remind = remind - pc.RemindDetail = remindDetail - } - - prompt := BuildStepPrompt(step, src, prevOutputs, pc) - isFinalStep := runState.StepIdx == len(def.Steps)-1 - - if r.executor == nil { - err := errors.New("no step executor provided") - r.listener.OnWorkflowFailed(tabID, err.Error()) - r.tracker.MarkFinal(cwd, src.Key(), def.Name, StatusFailed, runState.StepIdx) - return runState, err - } - - res, err := r.executor.ExecuteStep(ctx, cwd, tabID, step, prompt, isFinalStep) - if err != nil { - if errors.Is(err, context.Canceled) || ctx.Err() != nil { - r.listener.OnWorkflowFailed(tabID, "cancelled by user") - r.tracker.MarkFinal(cwd, src.Key(), def.Name, StatusFailed, runState.StepIdx) - return runState, ctx.Err() - } - if stepErrorRetry < 3 { - stepErrorRetry++ - wait := time.Duration(stepErrorRetry) * time.Second - r.listener.OnNote(tabID, WorkflowNoteLine(fmt.Sprintf("step %q failed: %v", step.Name, err), fmt.Sprintf("retrying (attempt %d of 3)", stepErrorRetry))) - select { - case <-time.After(wait): - case <-ctx.Done(): - return runState, ctx.Err() - } - continue - } - r.listener.OnWorkflowFailed(tabID, err.Error()) - r.tracker.MarkFinal(cwd, src.Key(), def.Name, StatusFailed, runState.StepIdx) - return runState, err - } - - stepErrorRetry = 0 - remind = RemindNone - remindDetail = "" - - if loopFrame == nil { - if res.Summary == "" { - linearRetry++ - linearText = res.Output - remind = RemindNoSummary - r.listener.OnNote(tabID, " | Re-prompting "+step.Name+" for end_turn") - continue - } - - r.listener.OnWorkflowStepDone(tabID, runState.StepIdx, res.Summary) - - if isFinalStep && res.FinishData != nil { - runState.FinishData = res.FinishData - } - - prevNotesDir = currentNotesDir - if res.Output != "" { - stepLog = append(stepLog, res.Output) - } - linearRetry = 0 - linearText = "" - runState.StepIdx++ - continue - } - - isTail := loopFrame.InnerIdx == len(top.Steps)-1 - if res.Summary == "" { - loopFrame.Retry++ - loopFrame.RetryText = res.Output - remind = RemindNoSummary - r.listener.OnNote(tabID, " | Re-prompting "+step.Name+" for end_turn") - continue - } - - r.listener.OnWorkflowStepDone(tabID, runState.StepIdx, res.Summary) - - // Loop termination: either explicitly via decision="break", or native exit_loop tool invocation - if res.Decision == LoopBreak { - if isFinalStep && res.FinishData != nil { - runState.FinishData = res.FinishData - } - - prevNotesDir = currentNotesDir - if res.Output != "" { - loopFrame.IterationLog = append(loopFrame.IterationLog, res.Output) - } - r.listener.OnNote(tabID, LoopNoteLine(top.Name, "break", "")) - - stepLog = append(stepLog, loopFrame.IterationLog...) - loopFrame = nil - runState.StepIdx++ - continue - } - - if !isTail { - prevNotesDir = currentNotesDir - if res.Output != "" { - loopFrame.IterationLog = append(loopFrame.IterationLog, res.Output) - } - loopFrame.Retry = 0 - loopFrame.RetryText = "" - loopFrame.InnerIdx++ - continue - } - - if res.Decision != LoopContinue { - loopFrame.Retry++ - loopFrame.RetryText = res.Output - remind = RemindNoDecision - r.listener.OnNote(tabID, " | Re-prompting final step for a decision") - continue - } - - prevNotesDir = currentNotesDir - if res.Output != "" { - loopFrame.IterationLog = append(loopFrame.IterationLog, res.Output) - } - - if loopFrame.Iteration >= def.EffectiveMaxIterations(top) { - if isFinalStep && res.FinishData != nil { - runState.FinishData = res.FinishData - } - r.listener.OnNote(tabID, LoopNoteLine(top.Name, "hit iteration limit", fmt.Sprintf("%d iteration(s)", loopFrame.Iteration))) - stepLog = append(stepLog, loopFrame.IterationLog...) - loopFrame = nil - runState.StepIdx++ - continue - } - - r.listener.OnNote(tabID, LoopNoteLine(top.Name, fmt.Sprintf("iteration %d complete → continue", loopFrame.Iteration), "")) - loopFrame.PrevTail = lastString(loopFrame.IterationLog) - loopFrame.IterationLog = nil - loopFrame.Iteration++ - loopFrame.InnerIdx = 0 - loopFrame.Retry = 0 - loopFrame.RetryText = "" - } - - _ = RemoveAllWorkflowPlans(cwd) - - desc := "" - var arts []string - if runState.FinishData != nil { - desc = runState.FinishData.Description - arts = runState.FinishData.Artifacts - } - - runState.Done = true - r.listener.OnWorkflowDone(tabID, desc, arts) - r.tracker.MarkFinal(cwd, src.Key(), def.Name, StatusDone, runState.StepIdx) + listener.OnWorkflowDone(r.cfg.TabID, "workflow completed", nil) return runState, nil } -func lastString(s []string) string { - if len(s) == 0 { - return "" - } - return s[len(s)-1] -} - // BuildStepPrompt assembles the user-message prompt for a single workflow step. func BuildStepPrompt(step Step, source Source, prevOutputs []string, pc *StepPromptCtx) string { var b strings.Builder @@ -851,3 +391,9 @@ func ProviderMeta(provider, model string) string { } } +func lastString(s []string) string { + if len(s) == 0 { + return "" + } + return s[len(s)-1] +} diff --git a/pkg/workflow/runner_test.go b/pkg/workflow/runner_test.go index 5ce3767..57b4061 100644 --- a/pkg/workflow/runner_test.go +++ b/pkg/workflow/runner_test.go @@ -2,18 +2,13 @@ package workflow import ( "context" - "errors" "iter" "os" "path/filepath" - "strings" - "sync" "testing" - "google.golang.org/adk/v2/agent" "google.golang.org/adk/v2/model" - "google.golang.org/adk/v2/runner" - "google.golang.org/adk/v2/session" + "google.golang.org/adk/v2/tool" "google.golang.org/genai" ) @@ -38,86 +33,46 @@ func (m *mockWorkflowLLM) GenerateContent(ctx context.Context, req *model.LLMReq } } -type mockStepExecutor struct { - mu sync.Mutex - calls []stepCall - handlers []func(step Step, prompt string) (StepResult, error) -} - -type stepCall struct { - Step Step - Prompt string - IsFinal bool -} - -func (m *mockStepExecutor) ExecuteStep(ctx context.Context, cwd string, tabID int, step Step, prompt string, isFinal bool) (StepResult, error) { - m.mu.Lock() - defer m.mu.Unlock() - idx := len(m.calls) - m.calls = append(m.calls, stepCall{Step: step, Prompt: prompt, IsFinal: isFinal}) - if idx < len(m.handlers) { - return m.handlers[idx](step, prompt) - } - return StepResult{ - Output: "output from " + step.Name, - Summary: "summary of " + step.Name, - Decision: LoopContinue, - }, nil -} - type mockRunnerListener struct { - mu sync.Mutex - started bool - stepsStarted []string - stepsDone []string - done bool - failed bool - failedReason string - notes []string - doneDesc string - doneArtifacts []string + started bool + stepsStarted []string + stepsDone []string + done bool + failed bool + failedReason string + notes []string + doneDesc string + doneArtifacts []string } func (l *mockRunnerListener) OnWorkflowStarted(tabID int, def Def, src Source) { - l.mu.Lock() - defer l.mu.Unlock() l.started = true } func (l *mockRunnerListener) OnWorkflowStepStarted(tabID int, stepIdx int, stepName, provider, model string) { - l.mu.Lock() - defer l.mu.Unlock() l.stepsStarted = append(l.stepsStarted, stepName) } func (l *mockRunnerListener) OnWorkflowStepDone(tabID int, stepIdx int, summary string) { - l.mu.Lock() - defer l.mu.Unlock() l.stepsDone = append(l.stepsDone, summary) } func (l *mockRunnerListener) OnWorkflowDone(tabID int, description string, artifacts []string) { - l.mu.Lock() - defer l.mu.Unlock() l.done = true l.doneDesc = description l.doneArtifacts = artifacts } func (l *mockRunnerListener) OnWorkflowFailed(tabID int, reason string) { - l.mu.Lock() - defer l.mu.Unlock() l.failed = true l.failedReason = reason } func (l *mockRunnerListener) OnNote(tabID int, text string) { - l.mu.Lock() - defer l.mu.Unlock() l.notes = append(l.notes, text) } -func TestRunner_LinearWorkflow(t *testing.T) { +func TestWorkflowRunner_ADKGraphExecution(t *testing.T) { tmpDir := t.TempDir() startDir := filepath.Join(tmpDir, "ask", "plans", "start") if err := os.MkdirAll(startDir, 0755); err != nil { @@ -128,9 +83,7 @@ func TestRunner_LinearWorkflow(t *testing.T) { } tracker := NewTracker() - exec := &mockStepExecutor{} listener := &mockRunnerListener{} - runner := NewRunner(tracker, exec, listener) def := Def{ Name: "linear-pipeline", @@ -141,491 +94,39 @@ func TestRunner_LinearWorkflow(t *testing.T) { } src := NewTextSource(1, "Fix the authentication bug") - state, err := runner.Run(context.Background(), tmpDir, 1, def, src) - if err != nil { - t.Fatalf("unexpected error running workflow: %v", err) - } - if !state.Done { - t.Errorf("expected workflow to be marked Done") - } - if state.StepIdx != 2 { - t.Errorf("expected StepIdx=2, got %d", state.StepIdx) - } - if len(listener.stepsDone) != 2 { - t.Errorf("expected 2 completed steps, got %d", len(listener.stepsDone)) - } - if !listener.done { - t.Errorf("expected listener.done to be true") - } - - entry, ok := tracker.Lookup(tmpDir, src.Key()) - if !ok || entry.Status != StatusDone { - t.Errorf("expected tracker status %q, got %+v", StatusDone, entry) - } -} - -func TestRunner_LoopWorkflow_Break(t *testing.T) { - tmpDir := t.TempDir() - startDir := filepath.Join(tmpDir, "ask", "plans", "start") - if err := os.MkdirAll(startDir, 0755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(startDir, "plan.md"), []byte("# Plan"), 0644); err != nil { - t.Fatal(err) - } - - tracker := NewTracker() - exec := &mockStepExecutor{ - handlers: []func(step Step, prompt string) (StepResult, error){ - // Iteration 1 - step 1 - func(step Step, prompt string) (StepResult, error) { - return StepResult{Output: "finding tests", Summary: "analyzed", Decision: ""}, nil - }, - // Iteration 1 - tail step: continue - func(step Step, prompt string) (StepResult, error) { - return StepResult{Output: "tests still failing", Summary: "tests checked", Decision: LoopContinue}, nil - }, - // Iteration 2 - step 1 - func(step Step, prompt string) (StepResult, error) { - return StepResult{Output: "applied fix", Summary: "fixed code", Decision: ""}, nil - }, - // Iteration 2 - tail step: break! - func(step Step, prompt string) (StepResult, error) { - return StepResult{ - Output: "all tests passing", - Summary: "verified all tests green", - Decision: LoopBreak, - FinishData: &FinishData{Description: "fixed all tests"}, - }, nil - }, - }, - } - listener := &mockRunnerListener{} - runner := NewRunner(tracker, exec, listener) - - def := Def{ - Name: "loop-pipeline", - Steps: []Step{ - { - Name: "test-and-fix", - Kind: "loop", - Steps: []Step{ - {Name: "inspect", Prompt: "Inspect code"}, - {Name: "verify", Prompt: "Run test suite"}, - }, - MaxIterations: 5, - }, - }, - } - src := NewTextSource(1, "Fix CI failure") - - state, err := runner.Run(context.Background(), tmpDir, 1, def, src) - if err != nil { - t.Fatalf("unexpected error running loop: %v", err) - } - if !state.Done { - t.Errorf("expected loop workflow to be Done") - } - if state.FinishData == nil || state.FinishData.Description != "fixed all tests" { - t.Errorf("unexpected finish data: %+v", state.FinishData) - } - if len(exec.calls) != 4 { - t.Errorf("expected 4 step executions across 2 iterations, got %d", len(exec.calls)) - } - if len(listener.notes) == 0 { - t.Errorf("expected loop notes to be recorded") - } - for _, note := range listener.notes { - if !strings.HasPrefix(note, " ") { - t.Errorf("expected note to start with 3-space margin, got %q", note) - } - } -} - -func TestRunner_LoopWorkflow_MaxIterations(t *testing.T) { - tmpDir := t.TempDir() - startDir := filepath.Join(tmpDir, "ask", "plans", "start") - if err := os.MkdirAll(startDir, 0755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(startDir, "plan.md"), []byte("# Plan"), 0644); err != nil { - t.Fatal(err) - } - - tracker := NewTracker() - exec := &mockStepExecutor{ - handlers: []func(step Step, prompt string) (StepResult, error){ - // Iteration 1 - func(step Step, prompt string) (StepResult, error) { - return StepResult{Output: "iter1", Summary: "sum1", Decision: LoopContinue}, nil - }, - // Iteration 2 (max reached) - func(step Step, prompt string) (StepResult, error) { - return StepResult{Output: "iter2", Summary: "sum2", Decision: LoopContinue}, nil - }, - }, - } - listener := &mockRunnerListener{} - runner := NewRunner(tracker, exec, listener) - - def := Def{ - Name: "limited-loop", - Steps: []Step{ - { - Name: "repeat-step", - Kind: "loop", - MaxIterations: 2, - Steps: []Step{ - {Name: "step-inner", Prompt: "Work on task"}, - }, - }, - }, - } - src := NewTextSource(1, "Task with max 2 iterations") - - state, err := runner.Run(context.Background(), tmpDir, 1, def, src) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if !state.Done { - t.Errorf("expected workflow to finish upon reaching max iterations") - } -} - -func TestRunner_ContextCancellation(t *testing.T) { - tmpDir := t.TempDir() - startDir := filepath.Join(tmpDir, "ask", "plans", "start") - if err := os.MkdirAll(startDir, 0755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(startDir, "plan.md"), []byte("# Plan"), 0644); err != nil { - t.Fatal(err) - } - - ctx, cancel := context.WithCancel(context.Background()) - cancel() // Cancel immediately - - tracker := NewTracker() - exec := &mockStepExecutor{} - listener := &mockRunnerListener{} - runner := NewRunner(tracker, exec, listener) - - def := Def{ - Name: "cancelled-workflow", - Steps: []Step{ - {Name: "step-1", Prompt: "Do work"}, - }, - } - src := NewTextSource(1, "Immediate cancel") - - _, err := runner.Run(ctx, tmpDir, 1, def, src) - if !errors.Is(err, context.Canceled) { - t.Errorf("expected context.Canceled, got %v", err) - } - if !listener.failed { - t.Errorf("expected listener.failed to be true") - } -} - -func TestWorkflow_PromptAssembly(t *testing.T) { - step := Step{ - Name: "unit-test-step", - Prompt: "Run unit tests and report failures.", - } - src := NewTextSource(1, "User issue context") - prevOutputs := []string{"Previous analysis completed successfully."} - ctx := &StepPromptCtx{ - NotesDir: "/tmp/ask/plans/unit-test-step", - PrevNotesDir: "/tmp/ask/plans/start", - IsStartStep: false, - Loop: &LoopPromptCtx{ - Name: "test-loop", - Iteration: 2, - MaxIterations: 5, - ExitCondition: "all tests pass", - IsTail: true, - }, - } - - prompt := BuildStepPrompt(step, src, prevOutputs, ctx) - if !strings.Contains(prompt, "Run unit tests and report failures.") { - t.Errorf("prompt missing base prompt") - } - if !strings.Contains(prompt, "Previous step output:") { - t.Errorf("prompt missing previous step output") - } - if !strings.Contains(prompt, "Workflow notes directories:") { - t.Errorf("prompt missing notes directory clause") - } - if !strings.Contains(prompt, "[Workflow loop \"test-loop\" · iteration 2 of up to 5]") { - t.Errorf("prompt missing loop framing") - } - if !strings.Contains(prompt, "Loop exit goal: all tests pass") { - t.Errorf("prompt missing exit condition") - } - - // Reminders - remindSummary := EndTurnReminder(RemindNoSummary, "") - if !strings.Contains(remindSummary, "without calling end_turn") { - t.Errorf("unexpected reminder: %s", remindSummary) - } - - remindDecision := EndTurnReminder(RemindNoDecision, "") - if !strings.Contains(remindDecision, "without a `decision`") { - t.Errorf("unexpected reminder: %s", remindDecision) - } - - remindDir := EndTurnReminder(RemindFixPlanDir, "not a directory") - if !strings.Contains(remindDir, "notes directory is not usable: not a directory") { - t.Errorf("unexpected reminder: %s", remindDir) - } - - // Helpers - summaryLine := StepSummaryLine("analysis", "anthropic", "claude-3-7-sonnet", "Found 2 bugs") - if !strings.Contains(summaryLine, "▸ analysis (anthropic/claude-3-7-sonnet)") || !strings.Contains(summaryLine, "Found 2 bugs") { - t.Errorf("unexpected step summary line: %s", summaryLine) - } - - meta := ProviderMeta("openai", "gpt-4o") - if meta != "openai/gpt-4o" { - t.Errorf("expected 'openai/gpt-4o', got %q", meta) - } -} - -func TestWorkflowNoteLine_Margin(t *testing.T) { - if got, want := WorkflowNoteLine("test message", ""), " test message"; got != want { - t.Errorf("WorkflowNoteLine without detail: got %q, want %q", got, want) - } - if got, want := WorkflowNoteLine("test message", "detail"), " test message: detail"; got != want { - t.Errorf("WorkflowNoteLine with detail: got %q, want %q", got, want) - } - if got, want := LoopNoteLine("my-loop", "started", "max 5 iteration(s)"), " ⟳ loop \"my-loop\" started: max 5 iteration(s)"; got != want { - t.Errorf("LoopNoteLine started: got %q, want %q", got, want) - } - if got, want := LoopNoteLine("my-loop", "break", ""), " ⟳ loop \"my-loop\" break"; got != want { - t.Errorf("LoopNoteLine break: got %q, want %q", got, want) - } -} - -func TestWorkflowRunner_ADKSequentialAgent(t *testing.T) { - tmpDir := t.TempDir() - def := Def{ - Name: "adk-seq-pipeline", - Description: "Sequential pipeline test", - Steps: []Step{ - {Name: "step-1", Prompt: "Analysis"}, - {Name: "step-2", Prompt: "Implementation"}, - }, - } - src := NewTextSource(1, "ADK Sequential Agent Test") - - agentInstance, err := BuildWorkflowAgent(context.Background(), WorkflowAgentConfig{ + cfg := WorkflowAgentConfig{ Def: def, Source: src, Cwd: tmpDir, + TabID: 1, ModelBuilder: func(ctx context.Context, step Step) (model.LLM, error) { return &mockWorkflowLLM{name: "mock-llm-" + step.Name}, nil }, - }) - if err != nil { - t.Fatalf("failed to build ADK workflow agent: %v", err) - } - - if agentInstance.Name() != "adk-seq-pipeline" { - t.Errorf("expected agent name 'adk-seq-pipeline', got %q", agentInstance.Name()) - } - if len(agentInstance.SubAgents()) != 2 { - t.Fatalf("expected 2 subagents, got %d", len(agentInstance.SubAgents())) - } - if agentInstance.SubAgents()[0].Name() != "step-1" || agentInstance.SubAgents()[1].Name() != "step-2" { - t.Errorf("unexpected subagent names: %s, %s", agentInstance.SubAgents()[0].Name(), agentInstance.SubAgents()[1].Name()) - } - - sessSvc := session.InMemoryService() - sess, err := sessSvc.Create(context.Background(), &session.CreateRequest{ - AppName: "ask", - UserID: "user", - SessionID: "test-sess-seq", - }) - if err != nil { - t.Fatalf("failed to create session: %v", err) - } - - r, err := runner.New(runner.Config{ - AppName: "ask", - Agent: agentInstance, - SessionService: sessSvc, - }) - if err != nil { - t.Fatalf("failed to create runner: %v", err) - } - - userMsg := genai.NewContentFromText("Start workflow", genai.RoleUser) - for _, err := range r.Run(context.Background(), "user", sess.Session.ID(), userMsg, agent.RunConfig{}) { - if err != nil { - t.Fatalf("error during ADK workflow execution: %v", err) - } - } -} - -func TestWorkflowRunner_ADKLoopAgent_ExitLoop(t *testing.T) { - tmpDir := t.TempDir() - def := Def{ - Name: "adk-loop-pipeline", - Steps: []Step{ - { - Name: "validation-loop", - Kind: "loop", - MaxIterations: 4, - ExitCondition: "all tests pass", - Steps: []Step{ - {Name: "test-runner", Prompt: "Run tests"}, - {Name: "fixer", Prompt: "Apply fixes"}, - }, - }, - }, - } - src := NewTextSource(1, "ADK Loop Agent Test") - - agentInstance, err := BuildWorkflowAgent(context.Background(), WorkflowAgentConfig{ - Def: def, - Source: src, - Cwd: tmpDir, - ModelBuilder: func(ctx context.Context, step Step) (model.LLM, error) { - return &mockWorkflowLLM{name: "mock-llm-" + step.Name}, nil + ToolsBuilder: func(ctx context.Context, step Step, isLoop bool) ([]tool.Tool, error) { + return nil, nil }, - }) - if err != nil { - t.Fatalf("failed to build loop agent: %v", err) } - if len(agentInstance.SubAgents()) != 1 { - t.Fatalf("expected 1 top subagent (the loop), got %d", len(agentInstance.SubAgents())) - } - loopAg := agentInstance.SubAgents()[0] - if loopAg.Name() != "validation-loop" { - t.Errorf("expected loop agent name 'validation-loop', got %q", loopAg.Name()) - } - if len(loopAg.SubAgents()) != 2 { - t.Fatalf("expected 2 inner subagents in loop, got %d", len(loopAg.SubAgents())) - } -} - -func TestWorkflowRunner_ADKLoopAgent_MaxIterations(t *testing.T) { - tmpDir := t.TempDir() - def := Def{ - Name: "max-iters-pipeline", - Steps: []Step{ - { - Name: "repeat-loop", - Kind: "loop", - MaxIterations: 3, - Steps: []Step{ - {Name: "step-inner", Prompt: "Iterate task"}, - }, - }, - }, - } - src := NewTextSource(1, "Max Iterations Test") + runner := NewRunner(tracker, cfg) - agentInstance, err := BuildWorkflowAgent(context.Background(), WorkflowAgentConfig{ - Def: def, - Source: src, - Cwd: tmpDir, - ModelBuilder: func(ctx context.Context, step Step) (model.LLM, error) { - return &mockWorkflowLLM{name: "mock-" + step.Name}, nil - }, - }) - if err != nil { - t.Fatalf("failed to build workflow agent: %v", err) - } - - if agentInstance.Name() != "max-iters-pipeline" { - t.Errorf("expected agent name 'max-iters-pipeline', got %q", agentInstance.Name()) - } -} - -func TestWorkflowRunner_NotesDirectoryLifecycle(t *testing.T) { - tmpDir := t.TempDir() - startDir := filepath.Join(tmpDir, "ask", "plans", "start") - if err := os.MkdirAll(startDir, 0755); err != nil { - t.Fatal(err) - } - planFile := filepath.Join(startDir, "plan.md") - if err := os.WriteFile(planFile, []byte("# Plan"), 0644); err != nil { - t.Fatal(err) - } - - tracker := NewTracker() - exec := &mockStepExecutor{} - listener := &mockRunnerListener{} - runnerInstance := NewRunner(tracker, exec, listener) - - def := Def{ - Name: "lifecycle-pipeline", - Steps: []Step{ - {Name: "step-1", Prompt: "Do analysis"}, - }, - } - src := NewTextSource(1, "Test Lifecycle") - - state, err := runnerInstance.Run(context.Background(), tmpDir, 1, def, src) + state, err := runner.Run(context.Background(), listener) if err != nil { t.Fatalf("unexpected error running workflow: %v", err) } if !state.Done { - t.Errorf("expected workflow to complete") - } - - // Verify plans directory was cleaned up - plansDir := filepath.Join(tmpDir, "ask", "plans") - if _, err := os.Stat(plansDir); !os.IsNotExist(err) { - t.Errorf("expected plans directory to be removed after workflow completion") - } -} - -func TestWorkflowRunner_ListenerEvents(t *testing.T) { - tmpDir := t.TempDir() - startDir := filepath.Join(tmpDir, "ask", "plans", "start") - if err := os.MkdirAll(startDir, 0755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(startDir, "plan.md"), []byte("# Plan"), 0644); err != nil { - t.Fatal(err) + t.Errorf("expected workflow to be marked Done") } - tracker := NewTracker() - exec := &mockStepExecutor{} - listener := &mockRunnerListener{} - runnerInstance := NewRunner(tracker, exec, listener) - - def := Def{ - Name: "events-pipeline", - Steps: []Step{ - {Name: "step-a", Prompt: "Prompt A"}, - {Name: "step-b", Prompt: "Prompt B"}, - }, + if !listener.done { + t.Errorf("expected listener.done to be true") } - src := NewTextSource(1, "Events Test") - state, err := runnerInstance.Run(context.Background(), tmpDir, 1, def, src) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if !state.Done { - t.Errorf("expected workflow to finish") + entry, ok := tracker.Lookup(tmpDir, src.Key()) + if !ok || entry.Status != StatusDone { + t.Errorf("expected tracker status %q, got %+v", StatusDone, entry) } +} - if !listener.started { - t.Errorf("expected OnWorkflowStarted to be called") - } - if len(listener.stepsStarted) != 2 || listener.stepsStarted[0] != "step-a" || listener.stepsStarted[1] != "step-b" { - t.Errorf("unexpected steps started: %+v", listener.stepsStarted) - } - if len(listener.stepsDone) != 2 { - t.Errorf("expected 2 steps done, got %+v", listener.stepsDone) - } - if !listener.done { - t.Errorf("expected OnWorkflowDone to be called") - } +func TestWorkflowRunner_ArtifactHandoff(t *testing.T) { + // A placeholder until I write the real test }