Threads 9: Align Thread with AI SDK request semantics - #266
Conversation
|
Bugbot is not enabled for your account, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
Reviewer's GuideAligns Thread’s run lifecycle and tree behavior with AI SDK semantics by treating assistant messages as continuations rather than parents of new runs, tracking a run-local linear path, reconstructing tool ownership on restore, and tightening concurrency and regeneration rules, with tests and docs updated accordingly. Sequence diagram for assistant continuation vs new run in sendMessagesequenceDiagram
actor User
participant AbstractThread
participant RunRegistry
participant ThreadRunChat
User->>AbstractThread: sendMessage(input, options?)
alt input is empty and cursor is assistant
AbstractThread->>AbstractThread: continueAssistant({ follow, messageId, options })
AbstractThread->>RunRegistry: createRunForAssistant(messageId)
RunRegistry-->>AbstractThread: RunRecord
AbstractThread->>RunRegistry: assertHasCapacity(run.spec.parentMessageId)
AbstractThread->>ThreadRunChat: run.chat.start(options)
ThreadRunChat-->>AbstractThread: finished
else normal new run
AbstractThread->>RunRegistry: assertHasCapacity(parentMessageId)
AbstractThread->>AbstractThread: startRun({ follow, from, input, options })
AbstractThread->>ThreadRunChat: startRunRequest(spec, chat.start)
ThreadRunChat-->>AbstractThread: finished
end
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Hey - I've found 2 issues, and left some high level feedback:
- The new
findAssistantOwningParthelpers scan the entire tree on each tool approval/output when a run is missing; consider indexing tool/approval IDs on restore (similar toRunRegistrymaps) to avoid repeated full-tree scans on larger threads. - With
ThreadRunStatenow caching#messagesinstead of reading from the host each time, please double-check that tree mutations outside the run (e.g. other runs or manual edits) cannot desynchronize this cache from the canonical tree; if they can, it may be safer to derive the path from the tree when needed or explicitly document the invariants.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The new `findAssistantOwningPart` helpers scan the entire tree on each tool approval/output when a run is missing; consider indexing tool/approval IDs on restore (similar to `RunRegistry` maps) to avoid repeated full-tree scans on larger threads.
- With `ThreadRunState` now caching `#messages` instead of reading from the host each time, please double-check that tree mutations outside the run (e.g. other runs or manual edits) cannot desynchronize this cache from the canonical tree; if they can, it may be safer to derive the path from the tree when needed or explicitly document the invariants.
## Individual Comments
### Comment 1
<location path="packages/thread/test/thread.test.ts" line_range="728-737" />
<code_context>
);
});
+ test("reconstructs tool and approval ownership after restoring a tree", async () => {
+ const source = new Thread();
+ source.addMessage(user("user-1"), null);
+ source.addMessage(
+ {
+ id: "assistant-1",
+ parts: [
+ {
+ approval: { id: "approval-1" },
+ input: { value: 1 },
+ state: "approval-requested",
+ toolCallId: "tool-1",
+ toolName: "test-tool",
+ type: "dynamic-tool",
+ },
+ ],
+ role: "assistant",
+ },
+ "user-1",
+ );
+ source.setCursor("assistant-1");
+ const restored = new Thread();
+ restored.restore(source.getTreeSnapshot());
+
+ await restored.addToolApprovalResponse({
+ approved: true,
+ id: "approval-1",
+ });
+ await restored.addToolOutput({
+ output: "restored output",
+ tool: "test-tool",
+ toolCallId: "tool-1",
+ });
+
+ expect(restored.getMessage("assistant-1")?.parts).toContainEqual(
+ expect.objectContaining({
+ approval: expect.objectContaining({
</code_context>
<issue_to_address>
**suggestion (testing):** Add tests for error paths in restored tool/approval ownership (missing or duplicate ownership).
To fully cover the new ownership reconstruction logic, please also add tests for:
1. A restored tree where the approval/toolCall id is missing from all assistant messages, asserting that `addToolApprovalResponse`/`addToolOutput` throw the expected error.
2. A restored tree where the same approval/toolCall id appears in multiple assistant messages, asserting that the duplicate-owner case throws the configured error.
These will exercise the new guardrails and help catch future regressions in error handling around ownership reconstruction.
</issue_to_address>
### Comment 2
<location path="packages/thread/test/ai-sdk-run-chat.test.ts" line_range="76-79" />
<code_context>
- getRunPath = () =>
- this.tree.getPath(this.spec.messageId ?? this.spec.parentMessageId);
+ getMessagePath = (messageId: string | null) => this.tree.getPath(messageId);
updateRunPath = (messages: UIMessage[]) => {
this.tree.updatePath(messages);
</code_context>
<issue_to_address>
**suggestion (testing):** Add tests for the new `initialPathMessageId`/`messages` tracking and `messageId` propagation in `ThreadRunChat`.
The new host/spec behavior (`getMessagePath`, `initialPathMessageId`) and the internal `#messages` array in `ThreadRunState` are not yet covered by tests. I’d suggest adding:
1. A `ThreadRunState` test that initializes `messages` from `getMessagePath(initialPathMessageId)` and verifies `pushMessage`, `popMessage`, and `replaceMessage` correctly update the internal array, including the guard that only allows replacing the current response.
2. A `sendMessages` test that confirms `messageId` is omitted only when `spec.messageId` is `undefined` and `trigger === "submit-message"`, and preserved for other triggers (e.g. `regenerate-message`) or when `spec.messageId` is set.
3. An `onFinish` test that checks it calls `host.getMessagePath(spec.messageId ?? spec.parentMessageId)` and returns the correct linear path for both explicit `messageId` and parent-only runs.
These will ensure the updated request/path semantics are validated at the run chat layer and guard against regressions in `messageId` and path handling.
Suggested implementation:
```typescript
this.tree = new MessageTree({ messages: [userMessage] });
}
getMessagePath = (messageId: string | null) => this.tree.getPath(messageId);
updateRunPath = (messages: UIMessage[]) => {
this.tree.updatePath(messages);
};
function createSpec(overrides: Partial<ThreadRunSpec> = {}): ThreadRunSpec {
return {
id: "run-1",
initialPathMessageId: "user-1",
parentMessageId: "user-1",
siblingOrder: 0,
...overrides,
};
}
describe("ThreadRunState message path tracking", () => {
it("initializes internal messages from getMessagePath(initialPathMessageId)", () => {
const userMessage: UIMessage = {
id: "user-1",
role: "user",
content: [{ type: "text", text: "Hello" }],
};
const host = new (class {
tree: MessageTree<UIMessage>;
constructor() {
this.tree = new MessageTree({ messages: [userMessage] });
}
getMessagePath = (messageId: string | null) =>
this.tree.getPath(messageId);
updateRunPath = (messages: UIMessage[]) => {
this.tree.updatePath(messages);
};
})();
const spec = createSpec({ initialPathMessageId: "user-1" });
const state = new ThreadRunState({ host, spec });
// The internal #messages array should be initialized from getMessagePath(initialPathMessageId)
const initialPath = host.getMessagePath(spec.initialPathMessageId);
expect(state.messages).toEqual(initialPath);
});
it("pushMessage appends the new message and updates host path", () => {
const userMessage: UIMessage = {
id: "user-1",
role: "user",
content: [{ type: "text", text: "Hello" }],
};
const host = new (class {
tree: MessageTree<UIMessage>;
constructor() {
this.tree = new MessageTree({ messages: [userMessage] });
}
getMessagePath = (messageId: string | null) =>
this.tree.getPath(messageId);
updateRunPath = jest.fn((messages: UIMessage[]) => {
this.tree.updatePath(messages);
});
})();
const spec = createSpec({ initialPathMessageId: "user-1" });
const state = new ThreadRunState({ host, spec });
const responseMessage: UIMessage = {
id: "assistant-1",
role: "assistant",
content: [{ type: "text", text: "World" }],
};
state.pushMessage(responseMessage);
expect(state.messages[state.messages.length - 1]).toEqual(responseMessage);
expect(host.updateRunPath).toHaveBeenCalledWith(state.messages);
});
it("popMessage removes the last message and updates host path", () => {
const userMessage: UIMessage = {
id: "user-1",
role: "user",
content: [{ type: "text", text: "Hello" }],
};
const responseMessage: UIMessage = {
id: "assistant-1",
role: "assistant",
content: [{ type: "text", text: "World" }],
};
const host = new (class {
tree: MessageTree<UIMessage>;
constructor() {
this.tree = new MessageTree({ messages: [userMessage, responseMessage] });
}
getMessagePath = (messageId: string | null) =>
this.tree.getPath(messageId);
updateRunPath = jest.fn((messages: UIMessage[]) => {
this.tree.updatePath(messages);
});
})();
const spec = createSpec({ initialPathMessageId: "assistant-1" });
const state = new ThreadRunState({ host, spec });
// sanity check: we start with both messages
expect(state.messages.map((m) => m.id)).toEqual(["user-1", "assistant-1"]);
state.popMessage();
expect(state.messages.map((m) => m.id)).toEqual(["user-1"]);
expect(host.updateRunPath).toHaveBeenCalledWith(state.messages);
});
it("replaceMessage only replaces the current response message", () => {
const userMessage: UIMessage = {
id: "user-1",
role: "user",
content: [{ type: "text", text: "Hello" }],
};
const responseMessage: UIMessage = {
id: "assistant-1",
role: "assistant",
content: [{ type: "text", text: "World" }],
};
const host = new (class {
tree: MessageTree<UIMessage>;
constructor() {
this.tree = new MessageTree({ messages: [userMessage, responseMessage] });
}
getMessagePath = (messageId: string | null) =>
this.tree.getPath(messageId);
updateRunPath = jest.fn((messages: UIMessage[]) => {
this.tree.updatePath(messages);
});
})();
const spec = createSpec({ initialPathMessageId: "assistant-1" });
const state = new ThreadRunState({ host, spec });
const newResponse: UIMessage = {
id: "assistant-2",
role: "assistant",
content: [{ type: "text", text: "New World" }],
};
// Replacing the current response should succeed
state.replaceMessage(newResponse);
expect(state.messages.map((m) => m.id)).toEqual(["user-1", "assistant-2"]);
expect(host.updateRunPath).toHaveBeenCalledWith(state.messages);
// Attempting to replace a non-current response should be ignored/guarded
const invalidReplacement: UIMessage = {
id: "user-2",
role: "user",
content: [{ type: "text", text: "Should not replace" }],
};
state.replaceMessage(invalidReplacement);
// The assistant response should still be the last message
expect(state.messages.map((m) => m.id)).toEqual(["user-1", "assistant-2"]);
});
});
describe("ThreadRunChat sendMessages messageId propagation", () => {
it("omits messageId only when spec.messageId is undefined and trigger === 'submit-message'", async () => {
const host = {
getMessagePath: jest.fn(),
updateRunPath: jest.fn(),
};
const specWithoutMessageId = createSpec({ parentMessageId: "user-1", initialPathMessageId: "user-1" });
const specWithMessageId = createSpec({
messageId: "user-2",
parentMessageId: "user-1",
initialPathMessageId: "user-2",
});
const threadRunWithoutMessageId = new ThreadRunChat({ host, spec: specWithoutMessageId });
const threadRunWithMessageId = new ThreadRunChat({ host, spec: specWithMessageId });
const messages: UIMessage[] = [
{
id: "user-1",
role: "user",
content: [{ type: "text", text: "Hello" }],
},
];
// submit-message trigger with undefined spec.messageId should omit messageId
const submitResult = await threadRunWithoutMessageId.sendMessages(messages, {
trigger: "submit-message",
});
expect(submitResult.request.messages[0].id).toBeUndefined();
// regenerate-message trigger should preserve messageId even when spec.messageId is undefined
const regenerateResult = await threadRunWithoutMessageId.sendMessages(messages, {
trigger: "regenerate-message",
});
expect(regenerateResult.request.messages[0].id).toBe("user-1");
// when spec.messageId is set, messageId should always be preserved
const withMessageIdResult = await threadRunWithMessageId.sendMessages(messages, {
trigger: "submit-message",
});
expect(withMessageIdResult.request.messages[0].id).toBe("user-1");
});
it("preserves messageId for non-submit triggers", async () => {
const host = {
getMessagePath: jest.fn(),
updateRunPath: jest.fn(),
};
const spec = createSpec({ parentMessageId: "user-1", initialPathMessageId: "user-1" });
const threadRun = new ThreadRunChat({ host, spec });
const messages: UIMessage[] = [
{
id: "user-1",
role: "user",
content: [{ type: "text", text: "Hello" }],
},
];
const result = await threadRun.sendMessages(messages, {
trigger: "regenerate-message",
});
expect(result.request.messages[0].id).toBe("user-1");
});
});
describe("ThreadRunChat onFinish path resolution", () => {
it("calls host.getMessagePath(spec.messageId ?? spec.parentMessageId) with explicit messageId", async () => {
const userMessage: UIMessage = {
id: "user-1",
role: "user",
content: [{ type: "text", text: "Hello" }],
};
const responseMessage: UIMessage = {
id: "assistant-1",
role: "assistant",
content: [{ type: "text", text: "World" }],
};
const host = {
tree: new MessageTree<UIMessage>({ messages: [userMessage, responseMessage] }),
getMessagePath: jest.fn(function (this: any, messageId: string | null) {
return this.tree.getPath(messageId);
}),
updateRunPath: jest.fn(),
};
const spec = createSpec({
messageId: "assistant-1",
parentMessageId: "user-1",
initialPathMessageId: "assistant-1",
});
const threadRun = new ThreadRunChat({ host, spec });
const resultPath = await threadRun.onFinish();
expect(host.getMessagePath).toHaveBeenCalledWith("assistant-1");
expect(resultPath.map((m) => m.id)).toEqual(["user-1", "assistant-1"]);
});
it("calls host.getMessagePath(spec.messageId ?? spec.parentMessageId) with parent-only runs", async () => {
const userMessage: UIMessage = {
id: "user-1",
role: "user",
content: [{ type: "text", text: "Hello" }],
};
const host = {
tree: new MessageTree<UIMessage>({ messages: [userMessage] }),
getMessagePath: jest.fn(function (this: any, messageId: string | null) {
return this.tree.getPath(messageId);
}),
updateRunPath: jest.fn(),
};
const spec = createSpec({
messageId: undefined,
parentMessageId: "user-1",
initialPathMessageId: "user-1",
});
const threadRun = new ThreadRunChat({ host, spec });
const resultPath = await threadRun.onFinish();
expect(host.getMessagePath).toHaveBeenCalledWith("user-1");
expect(resultPath.map((m) => m.id)).toEqual(["user-1"]);
});
});
```
To fully integrate these tests, the following adjustments will likely be needed elsewhere in `ai-sdk-run-chat.test.ts`:
1. **Imports**:
- Ensure `ThreadRunState`, `ThreadRunChat`, `MessageTree`, and `UIMessage` are imported from their respective modules. If the file already imports them, avoid duplicate imports; otherwise add:
- `import { ThreadRunState } from "../src/ThreadRunState";` (or the correct path in your repo)
- `import { ThreadRunChat } from "../src/ThreadRunChat";`
- `import { MessageTree } from "../src/MessageTree";`
- `import type { UIMessage } from "../src/types";`
- Make sure `jest` globals (`jest.fn`, `expect`, `describe`, `it`) are available; if using `vitest`, replace `jest.fn` with `vi.fn` and adjust imports accordingly.
2. **Internal API alignment**:
- The tests assume `ThreadRunState` exposes an internal `messages` array and methods `pushMessage`, `popMessage`, and `replaceMessage`. If these are private or named differently, either:
- Expose them for testing (e.g. via a getter), or
- Update the tests to use the public API that drives those methods.
- The `sendMessages` tests assume `ThreadRunChat.sendMessages` returns an object with `request.messages`. If your actual API differs, adjust the assertions to match the real shape.
3. **Trigger semantics**:
- If your `trigger` type is an enum or union with different string values than `"submit-message"` / `"regenerate-message"`, update the triggers in the tests to the correct values.
4. **Host mock shape**:
- Align the `host` mock objects with the actual `ThreadRunChat`/`ThreadRunState` constructor expectations (e.g. if the host is a class instance rather than a plain object, use the same pattern as existing tests in the file).
5. **createSpec helper**:
- If `ThreadRunSpec` includes additional required properties beyond `id`, `initialPathMessageId`, `parentMessageId`, and `siblingOrder`, extend the `createSpec` helper to include defaults for them so the tests construct valid specs.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
There was a problem hiding this comment.
All reported issues were addressed across 7 files
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
9e196c9 to
be36256
Compare
Summary\n- align isolated threaded runs with AI SDK continuation, regeneration, reconnection, and restored tool ownership behavior\n- preserve tree topology while each run exposes a linear AI SDK request path\n- document the assistant-parent regeneration blocker in AI SDK\n\n## Verification\n- bun test packages/thread/test\n- bun run lint (packages/thread)\n- bun test:types\n\nPart of the Threads stack.
Summary by Sourcery
Align thread run lifecycle with AI SDK request semantics, including assistant continuation, regeneration, reconnection, and restored tool ownership behavior.
New Features:
Bug Fixes:
Enhancements:
Build:
Documentation:
Tests:
Summary by cubic
Aligns
Threadrun behavior with AI SDK semantics: assistant continuations stay on the same message, regeneration creates sibling responses (including roots), and reconnection resumes the existing assistant node across restored trees. Preserves tree topology with a linear run-local request path and hardens resume/continuation and restored tool ownership.New Features
sendMessagecontinues a selected or explicit assistant in-place; no child is created, overlapping continuations are rejected, and concurrency is enforced.regenerate({ messageId })uses AI SDK’s native trigger, truncates only the run-local path, inserts a sibling replacement (supports roots), and only follows when still on the target.onFinishreceives the completed run path even after navigation.restore, enablingaddToolApprovalResponseandaddToolOutput; errors on missing or duplicate owners.Refactors
initialPathMessageIdand run-local path management; host now providesgetMessagePath.assertHasCapacityand can find runs for approvals/tool calls and response messages.ThreadRunChataddsstartWithMessage,regenerateMessage, andrefreshPath; alignssendMessagesID handling andonFinishwith the message path.test:typesrunstscbefore build.Written for commit f8e22ed. Summary will update on new commits.
Stack