Skip to content

Commit 11cd4fb

Browse files
committed
core: extract session entry stepping logic into dedicated module
Move the step function from session-entry.ts to session-entry-stepper.ts and remove immer dependency. Add static fromEvent factory methods to Synthetic, Assistant, and Compaction classes for cleaner event-to-entry conversion.
1 parent 9c16bd1 commit 11cd4fb

3 files changed

Lines changed: 446 additions & 173 deletions

File tree

Lines changed: 253 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,253 @@
1+
import { castDraft, produce, type WritableDraft } from "immer"
2+
import { SessionEvent } from "./session-event"
3+
import { SessionEntry } from "./session-entry"
4+
5+
export type MemoryState = {
6+
entries: SessionEntry.Entry[]
7+
pending: SessionEntry.Entry[]
8+
}
9+
10+
export interface Adapter<Result> {
11+
readonly getCurrentAssistant: () => SessionEntry.Assistant | undefined
12+
readonly updateAssistant: (assistant: SessionEntry.Assistant) => void
13+
readonly appendEntry: (entry: SessionEntry.Entry) => void
14+
readonly appendPending: (entry: SessionEntry.Entry) => void
15+
readonly finish: () => Result
16+
}
17+
18+
export function memory(state: MemoryState): Adapter<MemoryState> {
19+
const activeAssistantIndex = () =>
20+
state.entries.findLastIndex((entry) => entry.type === "assistant" && !entry.time.completed)
21+
22+
return {
23+
getCurrentAssistant() {
24+
const index = activeAssistantIndex()
25+
if (index < 0) return
26+
const assistant = state.entries[index]
27+
return assistant?.type === "assistant" ? assistant : undefined
28+
},
29+
updateAssistant(assistant) {
30+
const index = activeAssistantIndex()
31+
if (index < 0) return
32+
const current = state.entries[index]
33+
if (current?.type !== "assistant") return
34+
state.entries[index] = assistant
35+
},
36+
appendEntry(entry) {
37+
state.entries.push(entry)
38+
},
39+
appendPending(entry) {
40+
state.pending.push(entry)
41+
},
42+
finish() {
43+
return state
44+
},
45+
}
46+
}
47+
48+
export function stepWith<Result>(adapter: Adapter<Result>, event: SessionEvent.Event): Result {
49+
const currentAssistant = adapter.getCurrentAssistant()
50+
type DraftAssistant = WritableDraft<SessionEntry.Assistant>
51+
type DraftTool = WritableDraft<SessionEntry.AssistantTool>
52+
type DraftText = WritableDraft<SessionEntry.AssistantText>
53+
type DraftReasoning = WritableDraft<SessionEntry.AssistantReasoning>
54+
55+
const latestTool = (assistant: DraftAssistant | undefined, callID?: string) =>
56+
assistant?.content.findLast(
57+
(item): item is DraftTool => item.type === "tool" && (callID === undefined || item.callID === callID),
58+
)
59+
60+
const latestText = (assistant: DraftAssistant | undefined) =>
61+
assistant?.content.findLast((item): item is DraftText => item.type === "text")
62+
63+
const latestReasoning = (assistant: DraftAssistant | undefined) =>
64+
assistant?.content.findLast((item): item is DraftReasoning => item.type === "reasoning")
65+
66+
SessionEvent.Event.match(event, {
67+
prompt: (event) => {
68+
const entry = SessionEntry.User.fromEvent(event)
69+
if (currentAssistant) {
70+
adapter.appendPending(entry)
71+
return
72+
}
73+
adapter.appendEntry(entry)
74+
},
75+
synthetic: (event) => {
76+
adapter.appendEntry(SessionEntry.Synthetic.fromEvent(event))
77+
},
78+
"step.started": (event) => {
79+
if (currentAssistant) {
80+
adapter.updateAssistant(
81+
produce(currentAssistant, (draft) => {
82+
draft.time.completed = event.timestamp
83+
}),
84+
)
85+
}
86+
adapter.appendEntry(SessionEntry.Assistant.fromEvent(event))
87+
},
88+
"step.ended": (event) => {
89+
if (currentAssistant) {
90+
adapter.updateAssistant(
91+
produce(currentAssistant, (draft) => {
92+
draft.time.completed = event.timestamp
93+
draft.cost = event.cost
94+
draft.tokens = event.tokens
95+
}),
96+
)
97+
}
98+
},
99+
"text.started": () => {
100+
if (currentAssistant) {
101+
adapter.updateAssistant(
102+
produce(currentAssistant, (draft) => {
103+
draft.content.push({
104+
type: "text",
105+
text: "",
106+
})
107+
}),
108+
)
109+
}
110+
},
111+
"text.delta": (event) => {
112+
if (currentAssistant) {
113+
adapter.updateAssistant(
114+
produce(currentAssistant, (draft) => {
115+
const match = latestText(draft)
116+
if (match) match.text += event.delta
117+
}),
118+
)
119+
}
120+
},
121+
"text.ended": () => {},
122+
"tool.input.started": (event) => {
123+
if (currentAssistant) {
124+
adapter.updateAssistant(
125+
produce(currentAssistant, (draft) => {
126+
draft.content.push({
127+
type: "tool",
128+
callID: event.callID,
129+
name: event.name,
130+
time: {
131+
created: event.timestamp,
132+
},
133+
state: {
134+
status: "pending",
135+
input: "",
136+
},
137+
})
138+
}),
139+
)
140+
}
141+
},
142+
"tool.input.delta": (event) => {
143+
if (currentAssistant) {
144+
adapter.updateAssistant(
145+
produce(currentAssistant, (draft) => {
146+
const match = latestTool(draft, event.callID)
147+
// oxlint-disable-next-line no-base-to-string -- event.delta is a Schema.String (runtime string)
148+
if (match && match.state.status === "pending") match.state.input += event.delta
149+
}),
150+
)
151+
}
152+
},
153+
"tool.input.ended": () => {},
154+
"tool.called": (event) => {
155+
if (currentAssistant) {
156+
adapter.updateAssistant(
157+
produce(currentAssistant, (draft) => {
158+
const match = latestTool(draft, event.callID)
159+
if (match) {
160+
match.time.ran = event.timestamp
161+
match.state = {
162+
status: "running",
163+
input: event.input,
164+
}
165+
}
166+
}),
167+
)
168+
}
169+
},
170+
"tool.success": (event) => {
171+
if (currentAssistant) {
172+
adapter.updateAssistant(
173+
produce(currentAssistant, (draft) => {
174+
const match = latestTool(draft, event.callID)
175+
if (match && match.state.status === "running") {
176+
match.state = {
177+
status: "completed",
178+
input: match.state.input,
179+
output: event.output ?? "",
180+
title: event.title,
181+
metadata: event.metadata ?? {},
182+
attachments: [...(event.attachments ?? [])],
183+
}
184+
}
185+
}),
186+
)
187+
}
188+
},
189+
"tool.error": (event) => {
190+
if (currentAssistant) {
191+
adapter.updateAssistant(
192+
produce(currentAssistant, (draft) => {
193+
const match = latestTool(draft, event.callID)
194+
if (match && match.state.status === "running") {
195+
match.state = {
196+
status: "error",
197+
error: event.error,
198+
input: match.state.input,
199+
metadata: event.metadata ?? {},
200+
}
201+
}
202+
}),
203+
)
204+
}
205+
},
206+
"reasoning.started": () => {
207+
if (currentAssistant) {
208+
adapter.updateAssistant(
209+
produce(currentAssistant, (draft) => {
210+
draft.content.push({
211+
type: "reasoning",
212+
text: "",
213+
})
214+
}),
215+
)
216+
}
217+
},
218+
"reasoning.delta": (event) => {
219+
if (currentAssistant) {
220+
adapter.updateAssistant(
221+
produce(currentAssistant, (draft) => {
222+
const match = latestReasoning(draft)
223+
if (match) match.text += event.delta
224+
}),
225+
)
226+
}
227+
},
228+
"reasoning.ended": (event) => {
229+
if (currentAssistant) {
230+
adapter.updateAssistant(
231+
produce(currentAssistant, (draft) => {
232+
const match = latestReasoning(draft)
233+
if (match) match.text = event.text
234+
}),
235+
)
236+
}
237+
},
238+
retried: () => {},
239+
compacted: (event) => {
240+
adapter.appendEntry(SessionEntry.Compaction.fromEvent(event))
241+
},
242+
})
243+
244+
return adapter.finish()
245+
}
246+
247+
export function step(old: MemoryState, event: SessionEvent.Event): MemoryState {
248+
return produce(old, (draft) => {
249+
stepWith(memory(draft as MemoryState), event)
250+
})
251+
}
252+
253+
export * as SessionEntryStepper from "./session-entry-stepper"

0 commit comments

Comments
 (0)