From d51ddd6130f2893b67c8697ad2b83b703ac5fec4 Mon Sep 17 00:00:00 2001 From: NIKHIL ACHALE <96230495+nikhilachale@users.noreply.github.com> Date: Mon, 6 Jul 2026 02:58:08 +0530 Subject: [PATCH] feat: enhance Kiro agent integration with custom AO agent prompts (#2444) * feat: enhance Kiro agent integration with system prompts and agent configuration * fix(kiro): clear stale model from hooks when configuration is removed --- backend/internal/adapters/agent/kiro/hooks.go | 55 +++- backend/internal/adapters/agent/kiro/kiro.go | 54 ++-- .../internal/adapters/agent/kiro/kiro_test.go | 300 +++++++++++++++++- backend/internal/ports/agent.go | 8 +- backend/internal/session_manager/manager.go | 21 +- 5 files changed, 399 insertions(+), 39 deletions(-) diff --git a/backend/internal/adapters/agent/kiro/hooks.go b/backend/internal/adapters/agent/kiro/hooks.go index 04dc96a01..018870950 100644 --- a/backend/internal/adapters/agent/kiro/hooks.go +++ b/backend/internal/adapters/agent/kiro/hooks.go @@ -27,13 +27,18 @@ const ( // skips duplicates and uninstall recognizes AO entries by prefix without an // embedded template to diff against. kiroHookCommandPrefix = "ao hooks kiro " + + kiroAgentName = "ao" + kiroAgentDescription = "Agent Orchestrator session instructions" ) // kiroHookFile is the on-disk shape of .kiro/agents/ao.json. It is used by // tests to decode the written file. Kiro hooks are a map of camelCase event // name to a flat array of {matcher?, command} entries. type kiroHookFile struct { - Hooks map[string][]kiroHookEntry `json:"hooks"` + Name string `json:"name"` + Prompt *string `json:"prompt"` + Hooks map[string][]kiroHookEntry `json:"hooks"` } type kiroHookEntry struct { @@ -97,7 +102,7 @@ func (p *Plugin) GetAgentHooks(ctx context.Context, cfg ports.WorkspaceHookConfi } } - if err := writeKiroHooks(hooksPath, topLevel, rawHooks); err != nil { + if err := writeKiroHooks(hooksPath, topLevel, rawHooks, cfg.SystemPrompt, cfg.Config); err != nil { return fmt.Errorf("kiro.GetAgentHooks: %w", err) } if err := hookutil.EnsureWorkspaceGitignore(filepath.Dir(hooksPath), kiroAgentFileName); err != nil { @@ -137,7 +142,7 @@ func (p *Plugin) UninstallHooks(ctx context.Context, workspacePath string) error } } - if err := writeKiroHooks(hooksPath, topLevel, rawHooks); err != nil { + if err := writeKiroHooks(hooksPath, topLevel, rawHooks, "", ports.AgentConfig{}); err != nil { return fmt.Errorf("kiro.UninstallHooks: %w", err) } return nil @@ -210,7 +215,11 @@ func readKiroHooks(hooksPath string) (topLevel, rawHooks map[string]json.RawMess // writeKiroHooks folds rawHooks back into topLevel and writes the file. An // empty hooks map drops the "hooks" key entirely. -func writeKiroHooks(hooksPath string, topLevel, rawHooks map[string]json.RawMessage) error { +func writeKiroHooks(hooksPath string, topLevel, rawHooks map[string]json.RawMessage, systemPrompt string, agentConfig ports.AgentConfig) error { + if err := setKiroAgentDefaults(topLevel, systemPrompt, agentConfig); err != nil { + return err + } + if len(rawHooks) == 0 { delete(topLevel, "hooks") } else { @@ -235,6 +244,44 @@ func writeKiroHooks(hooksPath string, topLevel, rawHooks map[string]json.RawMess return nil } +func setKiroAgentDefaults(topLevel map[string]json.RawMessage, systemPrompt string, agentConfig ports.AgentConfig) error { + defaults := map[string]any{ + "name": kiroAgentName, + "description": kiroAgentDescription, + "prompt": nil, + "mcpServers": map[string]any{}, + "tools": []string{"*"}, + "toolAliases": map[string]any{}, + "allowedTools": []any{}, + "resources": []any{}, + "toolsSettings": map[string]any{}, + "includeMcpJson": true, + } + if systemPrompt != "" { + defaults["prompt"] = systemPrompt + } + if model := strings.TrimSpace(agentConfig.Model); model != "" { + defaults["model"] = model + } else { + delete(topLevel, "model") + } + + for key, value := range defaults { + managedKey := key == "name" || key == "prompt" || key == "model" + if !managedKey { + if _, ok := topLevel[key]; ok { + continue + } + } + data, err := json.Marshal(value) + if err != nil { + return fmt.Errorf("encode agent %s: %w", key, err) + } + topLevel[key] = data + } + return nil +} + // groupKiroHooksByEvent groups the managed hook specs by their Kiro event so // each event's array is rewritten once. func groupKiroHooksByEvent() map[string][]kiroHookSpec { diff --git a/backend/internal/adapters/agent/kiro/kiro.go b/backend/internal/adapters/agent/kiro/kiro.go index 5f54d73f0..0e386efcd 100644 --- a/backend/internal/adapters/agent/kiro/kiro.go +++ b/backend/internal/adapters/agent/kiro/kiro.go @@ -3,9 +3,9 @@ // and reading hook-derived session info. // // Kiro is AWS's agentic coding assistant. Its terminal CLI ships as the -// `kiro-cli` binary and exposes a non-interactive ("headless") mode via -// `kiro-cli chat --no-interactive ""`, suitable for AO-driven worker -// sessions. See https://kiro.dev/docs/cli/headless/ and +// `kiro-cli` binary. AO launches Kiro with a workspace-local custom agent so +// both worker and orchestrator sessions can use Kiro's normal interactive +// approval flow. See https://kiro.dev/docs/cli/headless/ and // https://kiro.dev/docs/cli/reference/cli-commands/. // // Launch delivers the initial prompt as a positional argument after `--` so a @@ -26,6 +26,7 @@ import ( "github.com/aoagents/agent-orchestrator/backend/internal/adapters" "github.com/aoagents/agent-orchestrator/backend/internal/adapters/agent/agentbase" "github.com/aoagents/agent-orchestrator/backend/internal/adapters/agent/binaryutil" + "github.com/aoagents/agent-orchestrator/backend/internal/domain" "github.com/aoagents/agent-orchestrator/backend/internal/ports" ) @@ -58,29 +59,47 @@ func (p *Plugin) Manifest() adapters.Manifest { } } -// GetLaunchCommand builds the argv to start a new headless Kiro session: -// `kiro-cli chat --no-interactive [trust flags] -- `. +// GetLaunchCommand builds the argv to start a new Kiro session: +// `kiro-cli chat --agent ao [trust flags] [-- ]`. // // The prompt is passed as a positional argument after `--` so a leading "-" is -// not read as a flag. Kiro's --no-interactive mode requires a prompt argument. +// not read as a flag. Kiro runs interactively for both workers and orchestrators; +// standing instructions come from the generated custom agent. func (p *Plugin) GetLaunchCommand(ctx context.Context, cfg ports.LaunchConfig) (cmd []string, err error) { binary, err := p.kiroBinary(ctx) if err != nil { return nil, err } - cmd = []string{binary, "chat", "--no-interactive"} + cmd = []string{binary, "chat", "--agent", kiroAgentName} appendApprovalFlags(&cmd, cfg.Permissions) - if cfg.Prompt != "" { - cmd = append(cmd, "--", cfg.Prompt) + prompt := cfg.Prompt + if prompt != "" { + cmd = append(cmd, "--", prompt) } return cmd, nil } -// GetRestoreCommand rebuilds the argv that continues an existing Kiro session: -// `kiro-cli chat --no-interactive --resume-id [trust flags]`. +// GetPromptDeliveryStrategy reports how Kiro receives the initial task prompt. +// Orchestrator standing instructions are delivered through the generated +// custom-agent prompt, so no command or post-start prompt injection is needed +// there. +func (p *Plugin) GetPromptDeliveryStrategy(ctx context.Context, cfg ports.LaunchConfig) (ports.PromptDeliveryStrategy, error) { + if err := ctx.Err(); err != nil { + return "", err + } + if cfg.Prompt != "" { + return ports.PromptDeliveryInCommand, nil + } + if cfg.Kind == domain.KindOrchestrator { + return ports.PromptDeliveryCustomAgent, nil + } + return ports.PromptDeliveryInCommand, nil +} + +// GetRestoreCommand rebuilds the argv that continues an existing Kiro session. // ok is false when the hook-derived native session id has not landed yet, so // callers can fall back to fresh launch behavior. func (p *Plugin) GetRestoreCommand(ctx context.Context, cfg ports.RestoreConfig) (cmd []string, ok bool, err error) { @@ -97,8 +116,7 @@ func (p *Plugin) GetRestoreCommand(ctx context.Context, cfg ports.RestoreConfig) return nil, false, err } - cmd = make([]string, 0, 8) - cmd = append(cmd, binary, "chat", "--no-interactive", "--resume-id", agentSessionID) + cmd = []string{binary, "chat", "--agent", kiroAgentName, "--resume-id", agentSessionID} appendApprovalFlags(&cmd, cfg.Permissions) return cmd, true, nil } @@ -152,14 +170,14 @@ func (p *Plugin) kiroBinary(ctx context.Context) (string, error) { return binary, nil } -// appendApprovalFlags maps AO's 4 permission modes onto Kiro's tool-trust -// flags. Default emits no flag so Kiro defers to the user's own configuration -// (the interactive per-tool prompt). accept-edits grants the write-capable -// built-in tools; auto/bypass grant all tools. +// appendApprovalFlags maps AO's permission modes onto Kiro's tool-trust flags. +// Default emits no flag so Kiro uses its normal interactive approval flow. +// accept-edits grants the write-capable built-in tools; auto/bypass grant all +// tools. func appendApprovalFlags(cmd *[]string, permissions ports.PermissionMode) { switch ports.NormalizePermissionMode(permissions) { case ports.PermissionModeDefault: - // No flag: defer to the user's Kiro config / per-tool prompting. + // No flag: defer to Kiro's normal interactive approval flow. case ports.PermissionModeAcceptEdits: *cmd = append(*cmd, "--trust-tools=fs_read,fs_write") case ports.PermissionModeAuto: diff --git a/backend/internal/adapters/agent/kiro/kiro_test.go b/backend/internal/adapters/agent/kiro/kiro_test.go index 6e4f437fd..e891ed4ef 100644 --- a/backend/internal/adapters/agent/kiro/kiro_test.go +++ b/backend/internal/adapters/agent/kiro/kiro_test.go @@ -11,6 +11,7 @@ import ( "github.com/aoagents/agent-orchestrator/backend/internal/adapters" "github.com/aoagents/agent-orchestrator/backend/internal/adapters/agent/authprobe" + "github.com/aoagents/agent-orchestrator/backend/internal/domain" "github.com/aoagents/agent-orchestrator/backend/internal/ports" ) @@ -27,7 +28,7 @@ func TestManifestIDIsKiro(t *testing.T) { } } -func TestGetLaunchCommandBuildsHeadlessArgv(t *testing.T) { +func TestGetLaunchCommandBuildsInteractiveArgv(t *testing.T) { plugin := &Plugin{resolvedBinary: "kiro-cli"} cmd, err := plugin.GetLaunchCommand(context.Background(), ports.LaunchConfig{ @@ -39,7 +40,8 @@ func TestGetLaunchCommandBuildsHeadlessArgv(t *testing.T) { } want := []string{ - "kiro-cli", "chat", "--no-interactive", + "kiro-cli", "chat", + "--agent", "ao", "--trust-all-tools", "--", "-fix this", } @@ -48,6 +50,64 @@ func TestGetLaunchCommandBuildsHeadlessArgv(t *testing.T) { } } +func TestGetLaunchCommandOrchestratorUsesInteractiveAgent(t *testing.T) { + plugin := &Plugin{resolvedBinary: "kiro-cli"} + + cmd, err := plugin.GetLaunchCommand(context.Background(), ports.LaunchConfig{ + Kind: domain.KindOrchestrator, + SystemPrompt: "You are the human-facing coordinator.", + }) + if err != nil { + t.Fatal(err) + } + + want := []string{ + "kiro-cli", "chat", + "--agent", "ao", + } + if !reflect.DeepEqual(cmd, want) { + t.Fatalf("unexpected command\nwant: %#v\n got: %#v", want, cmd) + } +} + +func TestGetLaunchCommandPromptlessWorkerStaysInteractive(t *testing.T) { + plugin := &Plugin{resolvedBinary: "kiro-cli"} + + cmd, err := plugin.GetLaunchCommand(context.Background(), ports.LaunchConfig{}) + if err != nil { + t.Fatal(err) + } + + want := []string{ + "kiro-cli", "chat", + "--agent", "ao", + } + if !reflect.DeepEqual(cmd, want) { + t.Fatalf("unexpected command\nwant: %#v\n got: %#v", want, cmd) + } +} + +func TestGetLaunchCommandPromptTakesPrecedenceOverSystemPrompt(t *testing.T) { + plugin := &Plugin{resolvedBinary: "kiro-cli"} + + cmd, err := plugin.GetLaunchCommand(context.Background(), ports.LaunchConfig{ + Prompt: "fix the failing test", + SystemPrompt: "standing role instructions", + }) + if err != nil { + t.Fatal(err) + } + + want := []string{ + "kiro-cli", "chat", + "--agent", "ao", + "--", "fix the failing test", + } + if !reflect.DeepEqual(cmd, want) { + t.Fatalf("unexpected command\nwant: %#v\n got: %#v", want, cmd) + } +} + func TestGetLaunchCommandMapsApprovalModes(t *testing.T) { tests := []struct { name string @@ -124,6 +184,33 @@ func TestGetPromptDeliveryStrategyIsInCommand(t *testing.T) { } } +func TestGetPromptDeliveryStrategyOrchestratorUsesCustomAgent(t *testing.T) { + plugin := &Plugin{resolvedBinary: "kiro-cli"} + + got, err := plugin.GetPromptDeliveryStrategy(context.Background(), ports.LaunchConfig{Kind: domain.KindOrchestrator}) + if err != nil { + t.Fatal(err) + } + if got != ports.PromptDeliveryCustomAgent { + t.Fatalf("unexpected strategy: %q", got) + } +} + +func TestGetPromptDeliveryStrategyPromptedOrchestratorUsesCommand(t *testing.T) { + plugin := &Plugin{resolvedBinary: "kiro-cli"} + + got, err := plugin.GetPromptDeliveryStrategy(context.Background(), ports.LaunchConfig{ + Kind: domain.KindOrchestrator, + Prompt: "do this explicit task", + }) + if err != nil { + t.Fatal(err) + } + if got != ports.PromptDeliveryInCommand { + t.Fatalf("unexpected strategy: %q", got) + } +} + func TestGetConfigSpecHasNoCustomFieldsYet(t *testing.T) { plugin := &Plugin{} @@ -182,7 +269,7 @@ func TestGetAgentHooksInstallsKiroHooks(t *testing.T) { t.Fatal(err) } hooksPath := filepath.Join(hooksDir, kiroAgentFileName) - existing := `{"name":"ao","hooks":{"stop":[{"command":"custom stop hook"}]}}` + existing := `{"name":"stale","hooks":{"stop":[{"command":"custom stop hook"}]}}` if err := os.WriteFile(hooksPath, []byte(existing), 0o644); err != nil { t.Fatal(err) } @@ -190,6 +277,7 @@ func TestGetAgentHooksInstallsKiroHooks(t *testing.T) { cfg := ports.WorkspaceHookConfig{ DataDir: t.TempDir(), SessionID: "sess-1", + SystemPrompt: "standing AO instructions", WorkspacePath: workspace, } if err := plugin.GetAgentHooks(context.Background(), cfg); err != nil { @@ -204,13 +292,23 @@ func TestGetAgentHooksInstallsKiroHooks(t *testing.T) { if err != nil { t.Fatal(err) } - // The unmanaged top-level "name" key must be preserved. var topLevel map[string]json.RawMessage if err := json.Unmarshal(data, &topLevel); err != nil { t.Fatal(err) } - if _, ok := topLevel["name"]; !ok { - t.Fatalf("unmanaged top-level key 'name' was dropped: %s", data) + var name string + if err := json.Unmarshal(topLevel["name"], &name); err != nil { + t.Fatalf("decode name from %s: %v", data, err) + } + if name != kiroAgentName { + t.Fatalf("name = %q, want %q", name, kiroAgentName) + } + var prompt string + if err := json.Unmarshal(topLevel["prompt"], &prompt); err != nil { + t.Fatalf("decode prompt from %s: %v", data, err) + } + if prompt != "standing AO instructions" { + t.Fatalf("prompt = %q, want system prompt", prompt) } var config kiroHookFile @@ -232,6 +330,167 @@ func TestGetAgentHooksInstallsKiroHooks(t *testing.T) { } } +func TestGetAgentHooksCreatesNamedKiroAgentFile(t *testing.T) { + plugin := &Plugin{resolvedBinary: "kiro-cli"} + workspace := t.TempDir() + hooksPath := kiroAgentPath(workspace) + + cfg := ports.WorkspaceHookConfig{ + DataDir: t.TempDir(), + SessionID: "sess-1", + SystemPrompt: "exact orchestrator system prompt", + WorkspacePath: workspace, + } + if err := plugin.GetAgentHooks(context.Background(), cfg); err != nil { + t.Fatal(err) + } + + data, err := os.ReadFile(hooksPath) + if err != nil { + t.Fatal(err) + } + var topLevel map[string]json.RawMessage + if err := json.Unmarshal(data, &topLevel); err != nil { + t.Fatal(err) + } + var name string + if err := json.Unmarshal(topLevel["name"], &name); err != nil { + t.Fatalf("decode name from %s: %v", data, err) + } + if name != kiroAgentName { + t.Fatalf("name = %q, want %q", name, kiroAgentName) + } + var config kiroHookFile + if err := json.Unmarshal(data, &config); err != nil { + t.Fatal(err) + } + if config.Prompt == nil || *config.Prompt != "exact orchestrator system prompt" { + t.Fatalf("prompt = %#v, want exact system prompt", config.Prompt) + } +} + +func TestGetAgentHooksWritesConfiguredModel(t *testing.T) { + plugin := &Plugin{resolvedBinary: "kiro-cli"} + workspace := t.TempDir() + hooksPath := kiroAgentPath(workspace) + + cfg := ports.WorkspaceHookConfig{ + Config: ports.AgentConfig{Model: "claude-sonnet-4-5"}, + DataDir: t.TempDir(), + SessionID: "sess-1", + SystemPrompt: "standing AO instructions", + WorkspacePath: workspace, + } + if err := plugin.GetAgentHooks(context.Background(), cfg); err != nil { + t.Fatal(err) + } + + data, err := os.ReadFile(hooksPath) + if err != nil { + t.Fatal(err) + } + var topLevel map[string]json.RawMessage + if err := json.Unmarshal(data, &topLevel); err != nil { + t.Fatal(err) + } + var model string + if err := json.Unmarshal(topLevel["model"], &model); err != nil { + t.Fatalf("decode model from %s: %v", data, err) + } + if model != "claude-sonnet-4-5" { + t.Fatalf("model = %q, want configured model", model) + } +} + +func TestGetAgentHooksOverwritesStaleConfiguredModel(t *testing.T) { + plugin := &Plugin{resolvedBinary: "kiro-cli"} + workspace := t.TempDir() + hooksPath := kiroAgentPath(workspace) + if err := os.MkdirAll(filepath.Dir(hooksPath), 0o755); err != nil { + t.Fatal(err) + } + existing := `{"name":"ao","model":"stale-model","tools":["custom"]}` + if err := os.WriteFile(hooksPath, []byte(existing), 0o644); err != nil { + t.Fatal(err) + } + + cfg := ports.WorkspaceHookConfig{ + Config: ports.AgentConfig{Model: "project-model"}, + DataDir: t.TempDir(), + SessionID: "sess-1", + SystemPrompt: "standing AO instructions", + WorkspacePath: workspace, + } + if err := plugin.GetAgentHooks(context.Background(), cfg); err != nil { + t.Fatal(err) + } + + data, err := os.ReadFile(hooksPath) + if err != nil { + t.Fatal(err) + } + var topLevel map[string]json.RawMessage + if err := json.Unmarshal(data, &topLevel); err != nil { + t.Fatal(err) + } + var model string + if err := json.Unmarshal(topLevel["model"], &model); err != nil { + t.Fatalf("decode model from %s: %v", data, err) + } + if model != "project-model" { + t.Fatalf("model = %q, want project override", model) + } + var tools []string + if err := json.Unmarshal(topLevel["tools"], &tools); err != nil { + t.Fatalf("decode tools from %s: %v", data, err) + } + if !reflect.DeepEqual(tools, []string{"custom"}) { + t.Fatalf("tools = %#v, want preserved custom tools", tools) + } +} + +func TestGetAgentHooksClearsStaleModelWhenConfigRemoved(t *testing.T) { + plugin := &Plugin{resolvedBinary: "kiro-cli"} + workspace := t.TempDir() + hooksPath := kiroAgentPath(workspace) + if err := os.MkdirAll(filepath.Dir(hooksPath), 0o755); err != nil { + t.Fatal(err) + } + existing := `{"name":"ao","model":"stale-model","tools":["custom"]}` + if err := os.WriteFile(hooksPath, []byte(existing), 0o644); err != nil { + t.Fatal(err) + } + + cfg := ports.WorkspaceHookConfig{ + DataDir: t.TempDir(), + SessionID: "sess-1", + SystemPrompt: "standing AO instructions", + WorkspacePath: workspace, + } + if err := plugin.GetAgentHooks(context.Background(), cfg); err != nil { + t.Fatal(err) + } + + data, err := os.ReadFile(hooksPath) + if err != nil { + t.Fatal(err) + } + var topLevel map[string]json.RawMessage + if err := json.Unmarshal(data, &topLevel); err != nil { + t.Fatal(err) + } + if _, ok := topLevel["model"]; ok { + t.Fatalf("model still present in %s, want cleared when config has no model", data) + } + var tools []string + if err := json.Unmarshal(topLevel["tools"], &tools); err != nil { + t.Fatalf("decode tools from %s: %v", data, err) + } + if !reflect.DeepEqual(tools, []string{"custom"}) { + t.Fatalf("tools = %#v, want preserved custom tools", tools) + } +} + func TestUninstallHooksRemovesKiroHooks(t *testing.T) { plugin := &Plugin{resolvedBinary: "kiro-cli"} workspace := t.TempDir() @@ -309,7 +568,8 @@ func TestGetRestoreCommandReadsAgentSessionID(t *testing.T) { t.Fatal("ok = false, want true") } want := []string{ - "kiro-cli", "chat", "--no-interactive", + "kiro-cli", "chat", + "--agent", "ao", "--resume-id", "uuid-123", "--trust-all-tools", } @@ -318,6 +578,32 @@ func TestGetRestoreCommandReadsAgentSessionID(t *testing.T) { } } +func TestGetRestoreCommandOrchestratorUsesInteractiveAgent(t *testing.T) { + plugin := &Plugin{resolvedBinary: "kiro-cli"} + + cmd, ok, err := plugin.GetRestoreCommand(context.Background(), ports.RestoreConfig{ + Kind: domain.KindOrchestrator, + Permissions: ports.PermissionModeDefault, + Session: ports.SessionRef{ + Metadata: map[string]string{ports.MetadataKeyAgentSessionID: "uuid-123"}, + }, + }) + if err != nil { + t.Fatalf("err = %v, want nil", err) + } + if !ok { + t.Fatal("ok = false, want true") + } + want := []string{ + "kiro-cli", "chat", + "--agent", "ao", + "--resume-id", "uuid-123", + } + if !reflect.DeepEqual(cmd, want) { + t.Fatalf("restore cmd\nwant: %#v\n got: %#v", want, cmd) + } +} + func TestGetRestoreCommandFalseWithoutAgentSessionID(t *testing.T) { plugin := &Plugin{resolvedBinary: "kiro-cli"} diff --git a/backend/internal/ports/agent.go b/backend/internal/ports/agent.go index 00859b7eb..5f1b342ca 100644 --- a/backend/internal/ports/agent.go +++ b/backend/internal/ports/agent.go @@ -130,6 +130,7 @@ const ( type LaunchConfig struct { Config AgentConfig IssueID string + Kind domain.SessionKind Permissions PermissionMode Prompt string SessionID string @@ -151,12 +152,14 @@ type WorkspaceHookConfig struct { Config AgentConfig DataDir string SessionID string + SystemPrompt string WorkspacePath string } // RestoreConfig carries inputs needed to continue an existing native agent session. type RestoreConfig struct { Config AgentConfig + Kind domain.SessionKind Permissions PermissionMode Session SessionRef // SystemPrompt carries the session's standing instructions (e.g. the @@ -217,6 +220,7 @@ type PromptDeliveryStrategy string // How the orchestrator hands the initial prompt to a freshly launched agent. const ( - PromptDeliveryInCommand PromptDeliveryStrategy = "in_command" - PromptDeliveryAfterStart PromptDeliveryStrategy = "after_start" + PromptDeliveryInCommand PromptDeliveryStrategy = "in_command" + PromptDeliveryAfterStart PromptDeliveryStrategy = "after_start" + PromptDeliveryCustomAgent PromptDeliveryStrategy = "custom_agent" ) diff --git a/backend/internal/session_manager/manager.go b/backend/internal/session_manager/manager.go index 11da9edd8..d55da460c 100644 --- a/backend/internal/session_manager/manager.go +++ b/backend/internal/session_manager/manager.go @@ -256,15 +256,16 @@ func (m *Manager) Spawn(ctx context.Context, cfg ports.SpawnConfig) (domain.Sess m.rollbackSpawnSeedRow(ctx, id) return domain.SessionRecord{}, fmt.Errorf("spawn %s: no agent adapter for harness %q", id, cfg.Harness) } - if err := m.prepareWorkspace(ctx, agent, id, ws.Path); err != nil { + agentConfig := effectiveAgentConfig(cfg.Kind, project.Config) + if err := m.prepareWorkspace(ctx, agent, id, ws.Path, systemPrompt, agentConfig); err != nil { _ = m.workspace.Destroy(ctx, ws) m.rollbackSpawnSeedRow(ctx, id) return domain.SessionRecord{}, fmt.Errorf("spawn %s: %w", id, err) } - agentConfig := effectiveAgentConfig(cfg.Kind, project.Config) argv, err := agent.GetLaunchCommand(ctx, ports.LaunchConfig{ SessionID: string(id), WorkspacePath: ws.Path, + Kind: cfg.Kind, Prompt: prompt, SystemPrompt: systemPrompt, IssueID: string(cfg.IssueID), @@ -575,9 +576,6 @@ func (m *Manager) Restore(ctx context.Context, id domain.SessionID) (domain.Sess if !ok { return domain.SessionRecord{}, fmt.Errorf("restore %s: no agent adapter for harness %q", id, rec.Harness) } - if err := m.prepareWorkspace(ctx, agent, id, ws.Path); err != nil { - return domain.SessionRecord{}, fmt.Errorf("restore %s: %w", id, err) - } // The system prompt is derived, not persisted: recompute it so a restored // session keeps its standing instructions across the relaunch. systemPrompt, err := m.buildSystemPrompt(ctx, rec.Kind, rec.ProjectID) @@ -586,7 +584,11 @@ func (m *Manager) Restore(ctx context.Context, id domain.SessionID) (domain.Sess } // Restore re-applies the project's resolved agent config so a configured // model/permissions carry across a restore, matching fresh spawn. - argv, err := restoreArgv(ctx, agent, id, ws.Path, meta, systemPrompt, effectiveAgentConfig(rec.Kind, project.Config), rec.Kind) + agentConfig := effectiveAgentConfig(rec.Kind, project.Config) + if err := m.prepareWorkspace(ctx, agent, id, ws.Path, systemPrompt, agentConfig); err != nil { + return domain.SessionRecord{}, fmt.Errorf("restore %s: %w", id, err) + } + argv, err := restoreArgv(ctx, agent, id, ws.Path, meta, systemPrompt, agentConfig, rec.Kind) if err != nil { return domain.SessionRecord{}, fmt.Errorf("restore %s: %w", id, err) } @@ -1297,11 +1299,13 @@ type preLauncher interface { // starts the agent: installing the workspace-local activity hooks (so early // startup hooks can update the already-created session row), then any optional // PreLaunch step. Shared by Spawn and Restore. -func (m *Manager) prepareWorkspace(ctx context.Context, agent ports.Agent, id domain.SessionID, workspacePath string) error { +func (m *Manager) prepareWorkspace(ctx context.Context, agent ports.Agent, id domain.SessionID, workspacePath, systemPrompt string, agentConfig ports.AgentConfig) error { if err := agent.GetAgentHooks(ctx, ports.WorkspaceHookConfig{ SessionID: string(id), WorkspacePath: workspacePath, DataDir: m.dataDir, + SystemPrompt: systemPrompt, + Config: agentConfig, }); err != nil { return fmt.Errorf("install hooks: %w", err) } @@ -1326,7 +1330,7 @@ func restoreArgv(ctx context.Context, agent ports.Agent, id domain.SessionID, wo WorkspacePath: workspacePath, Metadata: map[string]string{ports.MetadataKeyAgentSessionID: meta.AgentSessionID}, } - cmd, ok, err := agent.GetRestoreCommand(ctx, ports.RestoreConfig{Session: ref, SystemPrompt: systemPrompt, Config: agentConfig, Permissions: agentConfig.Permissions}) + cmd, ok, err := agent.GetRestoreCommand(ctx, ports.RestoreConfig{Session: ref, Kind: kind, SystemPrompt: systemPrompt, Config: agentConfig, Permissions: agentConfig.Permissions}) if err != nil { return nil, fmt.Errorf("restore command: %w", err) } @@ -1343,6 +1347,7 @@ func restoreArgv(ctx context.Context, agent ports.Agent, id domain.SessionID, wo argv, err := agent.GetLaunchCommand(ctx, ports.LaunchConfig{ SessionID: string(id), WorkspacePath: workspacePath, + Kind: kind, Prompt: meta.Prompt, SystemPrompt: systemPrompt, Config: agentConfig,