Skip to content

Commit 287aa89

Browse files
committed
fix(runtime): don't persist session_start hook output as a session message
session_start hook output (typically the AddEnvironmentInfo env block) was added to the session as a chat.MessageRoleSystem item via sess.AddMessage. Because session_start fires inside RunStream — after the caller has already appended the user's message — the system message landed AFTER the user's first turn. The user-message emitter at the top of RunStream reads messages[len-1] and surfaces it as the UserMessageEvent, so the env info leaked verbatim into the visible transcript on the user's first question. Fix: hold session_start AdditionalContext as transient extras for the duration of RunStream and thread it into sess.GetMessages alongside turn_start output on every iteration. The model still sees the env info on every model call; the persisted transcript and the user message tail stay clean. Same-shape contract as turn_start, which was already transient. - pkg/runtime/hooks.go: executeSessionStartHooks now returns []chat.Message instead of mutating the session. Both start-hook helpers share a small contextMessages converter so the bodies collapse to a single dispatchHook call each. - pkg/runtime/loop.go: capture session_start extras once at the top of RunStream; combine with per-iteration turn_start extras via slices.Concat when calling sess.GetMessages. - pkg/runtime/runtime_test.go: pin the regression with TestRunStream_AddEnvironmentInfo_DoesNotPolluteSession — asserts the session contains only user+assistant roles and that UserMessageEvent.Message is "hello", not the <env> block. Assisted-By: docker-agent
1 parent 579c933 commit 287aa89

3 files changed

Lines changed: 109 additions & 28 deletions

File tree

pkg/runtime/hooks.go

Lines changed: 27 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -85,32 +85,39 @@ func (r *LocalRuntime) dispatchHook(
8585
return result
8686
}
8787

88-
// executeSessionStartHooks executes session_start hooks and persists any
89-
// AdditionalContext as a system message on the session. SystemMessage,
90-
// if any, is emitted as a Warning by [dispatchHook].
91-
func (r *LocalRuntime) executeSessionStartHooks(ctx context.Context, sess *session.Session, a *agent.Agent, events chan Event) {
92-
result := r.dispatchHook(ctx, a, hooks.EventSessionStart, &hooks.Input{
88+
// executeSessionStartHooks fires session_start once at the top of
89+
// RunStream and returns its AdditionalContext as transient system
90+
// messages. The result is NOT persisted to the session: persisting
91+
// would pollute the visible transcript and (because session_start
92+
// fires after the user message has been added) shift the message the
93+
// runtime relays as the [UserMessageEvent]. Callers thread the
94+
// returned slice through [session.Session.GetMessages] on every
95+
// iteration so cwd / OS / arch context reaches the model without ever
96+
// being stored.
97+
func (r *LocalRuntime) executeSessionStartHooks(ctx context.Context, sess *session.Session, a *agent.Agent, events chan Event) []chat.Message {
98+
return contextMessages(r.dispatchHook(ctx, a, hooks.EventSessionStart, &hooks.Input{
9399
SessionID: sess.ID,
94100
Source: "startup",
95-
}, events)
96-
if result == nil || result.AdditionalContext == "" {
97-
return
98-
}
99-
slog.Debug("Session start hook provided additional context", "context", result.AdditionalContext)
100-
sess.AddMessage(session.SystemMessage(result.AdditionalContext))
101+
}, events))
101102
}
102103

103-
// executeTurnStartHooks runs turn_start hooks and returns ephemeral
104-
// system messages to inject into the model call's messages slice.
105-
//
106-
// Unlike session_start, the AdditionalContext from turn_start is NOT
107-
// persisted to the session — it's recomputed every turn. This is the
108-
// right semantics for fast-changing context like "Today's date" or the
109-
// contents of a prompt file the user might be editing during the session.
104+
// executeTurnStartHooks fires turn_start before each model call and
105+
// returns its AdditionalContext as transient system messages. Like
106+
// session_start the result is never persisted, but turn_start runs
107+
// every iteration so its content is recomputed each turn — the right
108+
// semantics for fast-changing context like the current date or the
109+
// contents of a prompt file the user might be editing mid-session.
110110
func (r *LocalRuntime) executeTurnStartHooks(ctx context.Context, sess *session.Session, a *agent.Agent, events chan Event) []chat.Message {
111-
result := r.dispatchHook(ctx, a, hooks.EventTurnStart, &hooks.Input{
111+
return contextMessages(r.dispatchHook(ctx, a, hooks.EventTurnStart, &hooks.Input{
112112
SessionID: sess.ID,
113-
}, events)
113+
}, events))
114+
}
115+
116+
// contextMessages converts a context-providing hook's AdditionalContext
117+
// into a one-element transient system-message slice ready to thread
118+
// through [session.Session.GetMessages]. Returns nil for empty results
119+
// so callers can pass it straight to [slices.Concat] without a guard.
120+
func contextMessages(result *hooks.Result) []chat.Message {
114121
if result == nil || result.AdditionalContext == "" {
115122
return nil
116123
}

pkg/runtime/loop.go

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -156,8 +156,12 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
156156

157157
a := r.resolveSessionAgent(sess)
158158

159-
// Execute session start hooks
160-
r.executeSessionStartHooks(ctx, sess, a, events)
159+
// session_start fires once per RunStream. Its AdditionalContext
160+
// (typically the AddEnvironmentInfo env block) is held as transient
161+
// extras and threaded into every model call below — never persisted,
162+
// to keep the visible transcript clean and the user message tail
163+
// stable.
164+
sessionStartMsgs := r.executeSessionStartHooks(ctx, sess, a, events)
161165

162166
// Emit team information
163167
events <- TeamInfo(r.agentDetailsFromTeam(), a.Name())
@@ -366,13 +370,14 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
366370
}
367371

368372
// Run turn_start hooks BEFORE building messages so their
369-
// AdditionalContext can be spliced after the invariant cache
370-
// checkpoint and before the conversation history. The hook
371-
// output is not persisted to the session, so per-turn signals
372-
// (date, prompt files) refresh every turn without bloating the
373-
// stored history.
373+
// AdditionalContext, alongside the session_start extras captured
374+
// once at the top of RunStream, can be spliced after the invariant
375+
// cache checkpoint and before the conversation history. Neither
376+
// hook's output is persisted, so per-turn signals (date, prompt
377+
// files) refresh every turn while session-level context (cwd, OS,
378+
// arch) stays stable — all without bloating the stored history.
374379
turnStartMsgs := r.executeTurnStartHooks(ctx, sess, a, events)
375-
messages := sess.GetMessages(a, turnStartMsgs...)
380+
messages := sess.GetMessages(a, slices.Concat(sessionStartMsgs, turnStartMsgs)...)
376381
slog.Debug("Retrieved messages for processing", "agent", a.Name(), "message_count", len(messages))
377382

378383
// Strip image content from messages if the model doesn't support image input.

pkg/runtime/runtime_test.go

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1994,6 +1994,75 @@ func TestRunStream_EmptyMessages_SendUserMessage(t *testing.T) {
19941994
require.NotEmpty(t, events)
19951995
}
19961996

1997+
// TestRunStream_AddEnvironmentInfo_DoesNotPolluteSession pins the
1998+
// regression where session_start hook output (the AddEnvironmentInfo
1999+
// env block) was persisted as a system message on the session AFTER
2000+
// the user's first message had already been added, then surfaced
2001+
// verbatim as the [UserMessageEvent] because the runtime relays
2002+
// messages[len-1] as the "current" user message.
2003+
func TestRunStream_AddEnvironmentInfo_DoesNotPolluteSession(t *testing.T) {
2004+
t.Parallel()
2005+
2006+
stream := newStreamBuilder().
2007+
AddContent("reply").
2008+
AddStopWithUsage(5, 5).
2009+
Build()
2010+
2011+
prov := &mockProvider{id: "test/mock-model", stream: stream}
2012+
root := agent.New(
2013+
"root", "You are a test agent",
2014+
agent.WithModel(prov),
2015+
agent.WithAddEnvironmentInfo(true),
2016+
)
2017+
tm := team.New(team.WithAgents(root))
2018+
2019+
rt, err := NewLocalRuntime(
2020+
tm,
2021+
WithSessionCompaction(false),
2022+
WithModelStore(mockModelStore{}),
2023+
WithWorkingDir(t.TempDir()),
2024+
)
2025+
require.NoError(t, err)
2026+
2027+
sess := session.New(
2028+
session.WithUserMessage("hello"),
2029+
session.WithWorkingDir(t.TempDir()),
2030+
)
2031+
2032+
evCh := rt.RunStream(t.Context(), sess)
2033+
var events []Event
2034+
for ev := range evCh {
2035+
events = append(events, ev)
2036+
}
2037+
2038+
// The persisted transcript must contain only the user message and
2039+
// the assistant reply — no system message smuggled in by the hook.
2040+
var roles []chat.MessageRole
2041+
for _, item := range sess.Messages {
2042+
if item.IsMessage() {
2043+
roles = append(roles, item.Message.Message.Role)
2044+
}
2045+
}
2046+
assert.Equal(t,
2047+
[]chat.MessageRole{chat.MessageRoleUser, chat.MessageRoleAssistant},
2048+
roles,
2049+
"session_start hook output must not be persisted as a session message",
2050+
)
2051+
2052+
// The UserMessageEvent must mirror the user's input, not the env
2053+
// info block produced by the hook.
2054+
var userEvts []*UserMessageEvent
2055+
for _, ev := range events {
2056+
if ue, ok := ev.(*UserMessageEvent); ok {
2057+
userEvts = append(userEvts, ue)
2058+
}
2059+
}
2060+
require.Len(t, userEvts, 1)
2061+
assert.Equal(t, "hello", userEvts[0].Message)
2062+
assert.NotContains(t, userEvts[0].Message, "<env>",
2063+
"user_message event must not leak the AddEnvironmentInfo block")
2064+
}
2065+
19972066
// recordingProvider wraps a sequence of mock streams and records the tools
19982067
// passed to each CreateChatCompletionStream call.
19992068
type recordingProvider struct {

0 commit comments

Comments
 (0)