diff --git a/pkg/agent/commands.go b/pkg/agent/commands.go index f16ee47..3ad5f6e 100644 --- a/pkg/agent/commands.go +++ b/pkg/agent/commands.go @@ -206,7 +206,8 @@ func (s *agent) handleCommand(ctx context.Context, input string, handler EventHa return } d := s.newDispatcher() - env := BuildAgentEnv(cfg, d, s.cwd) + newSessID := NewSessionID() + env := BuildAgentEnv(cfg, d, s.cwd, []string{"OLLIE_SESSION_ID=" + newSessID}) if env.Backend != "" { newBe, err := s.newBackend(env.Backend) if err != nil { @@ -229,7 +230,7 @@ func (s *agent) handleCommand(ctx context.Context, input string, handler EventHa s.cfg.MaxSteps = env.maxSteps s.agentName = name s.session = nil - s.sessionID = NewSessionID() + s.sessionID = newSessID s.pushSessionEnv() for _, msg := range env.Messages { handler(infoEvent(msg)) diff --git a/pkg/agent/core.go b/pkg/agent/core.go index f698694..9cb3a9c 100644 --- a/pkg/agent/core.go +++ b/pkg/agent/core.go @@ -47,8 +47,10 @@ type AgentEnv struct { // BuildAgentEnv constructs an AgentEnv from a pre-configured Dispatcher and // optional agent config. cwd sets the working directory reported in the // system prompt; if empty, the process working directory is used. +// env provides additional environment variables injected into prompt resolution +// subprocesses (e.g. OLLIE_SESSION_ID=xxx). // The caller is responsible for registering all servers on d before calling this. -func BuildAgentEnv(cfg *config.Config, d tools.Dispatcher, cwd string) AgentEnv { +func BuildAgentEnv(cfg *config.Config, d tools.Dispatcher, cwd string, env []string) AgentEnv { var messages []string var allToolInfos []tools.ToolInfo @@ -75,7 +77,7 @@ func BuildAgentEnv(cfg *config.Config, d tools.Dispatcher, cwd string) AgentEnv for k, v := range cfg.Hooks { hooks[k] = []string(v) } - if resolved, err := resolvePrompt(cfg.Prompt, cwd); err != nil { + if resolved, err := resolvePrompt(cfg.Prompt, cwd, env); err != nil { fmt.Fprintf(os.Stderr, "resolve prompt: %v\n", err) } else { preamble = resolved diff --git a/pkg/agent/core_test.go b/pkg/agent/core_test.go index 984306b..21b7ff2 100644 --- a/pkg/agent/core_test.go +++ b/pkg/agent/core_test.go @@ -2226,7 +2226,7 @@ func setupCfgDir(t *testing.T) string { func TestBuildAgentEnv_NilConfig(t *testing.T) { setupCfgDir(t) d := tools.NewDispatcher() - env := BuildAgentEnv(nil, d, t.TempDir()) + env := BuildAgentEnv(nil, d, t.TempDir(), nil) if env.preamble != "" { t.Errorf("preamble = %q; want empty", env.preamble) @@ -2246,7 +2246,7 @@ func TestBuildAgentEnv_PromptBecomesPreamble(t *testing.T) { setupCfgDir(t) d := tools.NewDispatcher() cfg := &config.Config{Prompt: config.Prompt{Value: []string{"the prompt"}}} - env := BuildAgentEnv(cfg, d, t.TempDir()) + env := BuildAgentEnv(cfg, d, t.TempDir(), nil) if env.preamble != "the prompt" { t.Errorf("preamble = %q; want %q", env.preamble, "the prompt") @@ -2262,7 +2262,7 @@ func TestBuildAgentEnv_HooksAndParams(t *testing.T) { MaxTokens: 512, Temperature: &temp, } - env := BuildAgentEnv(cfg, d, t.TempDir()) + env := BuildAgentEnv(cfg, d, t.TempDir(), nil) if cmds := env.Hooks[HookPreTurn]; len(cmds) != 1 || cmds[0] != "echo hi" { t.Errorf("Hooks[preTurn] = %v", cmds) @@ -2278,7 +2278,7 @@ func TestBuildAgentEnv_HooksAndParams(t *testing.T) { func TestBuildAgentEnv_ExecUnknownTool(t *testing.T) { setupCfgDir(t) d := tools.NewDispatcher() - env := BuildAgentEnv(nil, d, t.TempDir()) + env := BuildAgentEnv(nil, d, t.TempDir(), nil) _, err := env.exec(context.Background(), "no_such_tool", json.RawMessage(`{}`)) if err == nil || !strings.Contains(err.Error(), "unknown tool") { @@ -2318,7 +2318,7 @@ func (m *mockDispatcher) Dispatch(ctx context.Context, server, tool string, args func TestBuildAgentEnv_ListToolsError(t *testing.T) { setupCfgDir(t) d := &mockDispatcher{listErr: fmt.Errorf("boom")} - env := BuildAgentEnv(nil, d, t.TempDir()) + env := BuildAgentEnv(nil, d, t.TempDir(), nil) if len(env.Messages) != 1 || !strings.Contains(env.Messages[0], "boom") { t.Errorf("Messages = %v", env.Messages) } @@ -2329,7 +2329,7 @@ func TestBuildAgentEnv_ToolsPopulated(t *testing.T) { d := &mockDispatcher{ tools: []tools.ToolInfo{{Server: "s1", Name: "mytool", Description: "desc", InputSchema: json.RawMessage(`{}`)}}, } - env := BuildAgentEnv(nil, d, t.TempDir()) + env := BuildAgentEnv(nil, d, t.TempDir(), nil) if len(env.tools) != 1 || env.tools[0].Name != "mytool" { t.Errorf("tools = %+v", env.tools) } @@ -2342,7 +2342,7 @@ func TestBuildAgentEnv_ToolsDisabled(t *testing.T) { } f := false cfg := &config.Config{Tools: &f} - env := BuildAgentEnv(cfg, d, t.TempDir()) + env := BuildAgentEnv(cfg, d, t.TempDir(), nil) if len(env.tools) != 0 { t.Errorf("expected no tools when disabled; got %+v", env.tools) } @@ -2352,7 +2352,7 @@ func TestBuildAgentEnv_PromptOnly(t *testing.T) { setupCfgDir(t) d := tools.NewDispatcher() cfg := &config.Config{Prompt: config.Prompt{Value: []string{"only agent"}}} - env := BuildAgentEnv(cfg, d, t.TempDir()) + env := BuildAgentEnv(cfg, d, t.TempDir(), nil) if env.preamble != "only agent" { t.Errorf("preamble = %q; want %q", env.preamble, "only agent") } @@ -2365,7 +2365,7 @@ func TestBuildAgentEnv_ExecPrompt(t *testing.T) { Value: []string{"echo hello", "echo 'You are a security auditor.'", "echo world"}, IsExec: true, }} - env := BuildAgentEnv(cfg, d, t.TempDir()) + env := BuildAgentEnv(cfg, d, t.TempDir(), nil) want := "hello\nYou are a security auditor.\nworld" if env.preamble != want { t.Errorf("preamble = %q; want %q", env.preamble, want) @@ -2380,7 +2380,7 @@ func TestBuildAgentEnv_ExecDispatchSuccess(t *testing.T) { return json.RawMessage(`{"content":[{"type":"text","text":"ok"}]}`), nil }, } - env := BuildAgentEnv(nil, d, t.TempDir()) + env := BuildAgentEnv(nil, d, t.TempDir(), nil) result, err := env.exec(context.Background(), "mytool", json.RawMessage(`{}`)) if err != nil { t.Fatal(err) @@ -2398,7 +2398,7 @@ func TestBuildAgentEnv_ExecDispatchError(t *testing.T) { return nil, fmt.Errorf("dispatch failed") }, } - env := BuildAgentEnv(nil, d, t.TempDir()) + env := BuildAgentEnv(nil, d, t.TempDir(), nil) _, err := env.exec(context.Background(), "mytool", json.RawMessage(`{}`)) if err == nil || !strings.Contains(err.Error(), "dispatch failed") { t.Errorf("err = %v", err) @@ -2413,7 +2413,7 @@ func TestBuildAgentEnv_ExecToolResultIsError(t *testing.T) { return json.RawMessage(`{"isError":true,"content":[{"type":"text","text":"bad thing"}]}`), nil }, } - env := BuildAgentEnv(nil, d, t.TempDir()) + env := BuildAgentEnv(nil, d, t.TempDir(), nil) _, err := env.exec(context.Background(), "mytool", json.RawMessage(`{}`)) if err == nil || !strings.Contains(err.Error(), "bad thing") { t.Errorf("err = %v", err) diff --git a/pkg/agent/prompt_resolver.go b/pkg/agent/prompt_resolver.go index 990bfed..956900b 100644 --- a/pkg/agent/prompt_resolver.go +++ b/pkg/agent/prompt_resolver.go @@ -24,17 +24,17 @@ import ( // - If the string starts with '!', the rest is executed as a shell command. // - If the expanded string names an existing file, the file is read. // - Otherwise the string is used as-is. -func resolvePrompt(p config.Prompt, cwd string) (string, error) { +func resolvePrompt(p config.Prompt, cwd string, env []string) (string, error) { if len(p.Value) == 0 { return "", nil } if p.IsExec { - return resolveExecPrompt(p.Value, cwd) + return resolveExecPrompt(p.Value, cwd, env) } return resolveStringPrompt(p.Value[0], cwd) } -func resolveExecPrompt(cmds []string, cwd string) (string, error) { +func resolveExecPrompt(cmds []string, cwd string, env []string) (string, error) { var parts []string for _, cmdStr := range cmds { cmdStr = strings.TrimSpace(cmdStr) @@ -46,6 +46,9 @@ func resolveExecPrompt(cmds []string, cwd string) (string, error) { if cwd != "" { cmd.Dir = cwd } + if len(env) > 0 { + cmd.Env = append(os.Environ(), env...) + } var stdout, stderr bytes.Buffer cmd.Stdout = &stdout cmd.Stderr = &stderr