lkn
/
ollie-core
Archived
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: inject OLLIE_SESSION_ID into prompt resolution subprocesses 

resolveExecPrompt spawns x/prime which uses envsubst to expand
${OLLIE_SESSION_ID} in the system prompt template. The variable was
never in the process environment at resolution time — it was only
set on the execute server after the prompt was already rendered.

Add env []string parameter to resolvePrompt/resolveExecPrompt and
BuildAgentEnv so callers can inject session-scoped vars into the
subprocess environment. Generate session ID before BuildAgentEnv in
the agent-switch path so it is available at prompt resolution time.
This commit is contained in:
Levi Neely 2026-05-18 12:31:35 +02:00
parent 1207e26a07
commit 5efed7040d
4 changed files with 25 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
pkg/agent/commands.go
View File

@ -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))

6
pkg/agent/core.go
View File

@ -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

24
pkg/agent/core_test.go
View File

@ -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)

9
pkg/agent/prompt_resolver.go
View File

@ -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