Skip to content

Commit 2543cc1

Browse files
committed
chat refactor
Signed-off-by: Yujong Lee <yujonglee.dev@gmail.com>
1 parent c23b1f1 commit 2543cc1

File tree

10 files changed

+127
-138
lines changed

10 files changed

+127
-138
lines changed

apps/desktop/src/chat/components/session.tsx

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,10 @@ import { commands as templateCommands } from "@hypr/plugin-template";
1414

1515
import { useLanguageModel } from "~/ai/hooks";
1616
import type { ContextEntity, ContextRef } from "~/chat/context/entities";
17+
import { hydrateSessionContextFromFs } from "~/chat/context/session-context-hydrator";
1718
import { useChatContextPipeline } from "~/chat/context/use-chat-context-pipeline";
18-
import { hydrateSessionContextFromFs } from "~/chat/session-context-hydrator";
1919
import { useCreateChatMessage } from "~/chat/store/useCreateChatMessage";
20+
import { CONTEXT_TEXT_FIELD } from "~/chat/tools";
2021
import { CustomChatTransport } from "~/chat/transport";
2122
import type { HyprUIMessage } from "~/chat/types";
2223
import { useToolRegistry } from "~/contexts/tool";
@@ -63,7 +64,7 @@ function stripEphemeralToolContext(
6364
part.type !== "tool-search_sessions" ||
6465
part.state !== "output-available" ||
6566
!isRecord(part.output) ||
66-
!("contextText" in part.output)
67+
!(CONTEXT_TEXT_FIELD in part.output)
6768
) {
6869
return part;
6970
}
@@ -359,7 +360,7 @@ function useTransport(
359360
tools,
360361
effectiveSystemPrompt,
361362
async (ref) => {
362-
if (ref.kind !== "session" || !store) {
363+
if (!store) {
363364
return null;
364365
}
365366
return hydrateSessionContextFromFs(store, ref.sessionId);

apps/desktop/src/chat/context/entities.ts

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -90,19 +90,13 @@ function parseSearchSessionsOutput(output: unknown): ContextEntity[] {
9090
});
9191
}
9292

93-
export type ToolContextExtractor = (output: unknown) => ContextEntity[];
94-
95-
const toolEntityExtractors: Record<string, ToolContextExtractor> = {
93+
const toolEntityExtractors: Record<
94+
string,
95+
(output: unknown) => ContextEntity[]
96+
> = {
9697
search_sessions: parseSearchSessionsOutput,
9798
};
9899

99-
export function registerToolContextExtractor(
100-
toolName: string,
101-
extractor: ToolContextExtractor,
102-
): void {
103-
toolEntityExtractors[toolName] = extractor;
104-
}
105-
106100
export function extractToolContextEntities(
107101
messages: Array<Pick<HyprUIMessage, "parts">>,
108102
): ContextEntity[] {
File renamed without changes.

apps/desktop/src/chat/context/use-chat-context-pipeline.ts

Lines changed: 2 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -27,21 +27,12 @@ function getSessionDisplayData(
2727
};
2828
}
2929

30-
// Normalize legacy "session:current" key to the per-session key format.
31-
function normalizeRef(ref: ContextRef): ContextRef {
32-
if (ref.source === "auto-current" && ref.key === "session:current") {
33-
return { ...ref, key: `session:auto:${ref.sessionId}` };
34-
}
35-
return ref;
36-
}
37-
3830
function extractCommittedRefs(messages: HyprUIMessage[]): ContextRef[] {
3931
const seen = new Set<string>();
4032
const refs: ContextRef[] = [];
4133
for (const msg of messages) {
4234
if (msg.role !== "user") continue;
43-
for (const raw of msg.metadata?.contextRefs ?? []) {
44-
const ref = normalizeRef(raw);
35+
for (const ref of msg.metadata?.contextRefs ?? []) {
4536
if (!seen.has(ref.key)) {
4637
seen.add(ref.key);
4738
refs.push(ref);
@@ -88,9 +79,7 @@ export function useChatContextPipeline({
8879
sessionId: currentSessionId,
8980
});
9081
}
91-
for (const ref of pendingManualRefs) {
92-
refs.push(ref);
93-
}
82+
refs.push(...pendingManualRefs);
9483
return refs;
9584
}, [currentSessionId, pendingManualRefs]);
9685

apps/desktop/src/chat/tools/index.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@ import type { SearchFilters } from "~/search/contexts/engine/types";
1212

1313
export type { ToolDependencies };
1414

15+
// Ephemeral field: injected by transport during hydration, stripped before persistence.
16+
export const CONTEXT_TEXT_FIELD = "contextText" as const;
17+
1518
export const buildChatTools = (deps: ToolDependencies) => ({
1619
search_sessions: buildSearchSessionsTool(deps),
1720
search_contacts: buildSearchContactsTool(deps),
Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
import type { ContextRef } from "../context/entities";
2+
import { CONTEXT_TEXT_FIELD } from "../tools";
3+
4+
export function isRecord(value: unknown): value is Record<string, unknown> {
5+
return typeof value === "object" && value !== null;
6+
}
7+
8+
export const MAX_TOOL_STEPS = 5;
9+
export const MESSAGE_WINDOW_THRESHOLD = 20;
10+
export const MESSAGE_WINDOW_SIZE = 10;
11+
12+
export function isContextRef(value: unknown): value is ContextRef {
13+
return (
14+
isRecord(value) &&
15+
value.kind === "session" &&
16+
typeof value.key === "string" &&
17+
typeof value.sessionId === "string" &&
18+
(value.source === undefined ||
19+
value.source === "tool" ||
20+
value.source === "manual" ||
21+
value.source === "auto-current")
22+
);
23+
}
24+
25+
export function getContextRefs(metadata: unknown): ContextRef[] {
26+
if (!isRecord(metadata) || !Array.isArray(metadata.contextRefs)) {
27+
return [];
28+
}
29+
30+
return metadata.contextRefs.filter((ref): ref is ContextRef =>
31+
isContextRef(ref),
32+
);
33+
}
34+
35+
export function getSessionIdsFromSearchOutput(output: unknown): string[] {
36+
if (!isRecord(output) || !Array.isArray(output.results)) {
37+
return [];
38+
}
39+
return output.results.flatMap((item) => {
40+
if (
41+
!isRecord(item) ||
42+
(typeof item.id !== "string" && typeof item.id !== "number")
43+
) {
44+
return [];
45+
}
46+
return [String(item.id)];
47+
});
48+
}
49+
50+
export type ToolOutputPart = {
51+
type: `tool-${string}`;
52+
state: "output-available";
53+
output?: unknown;
54+
[key: string]: unknown;
55+
};
56+
57+
export function isToolOutputPart(value: unknown): value is ToolOutputPart {
58+
return (
59+
isRecord(value) &&
60+
typeof value.type === "string" &&
61+
value.type.startsWith("tool-") &&
62+
value.state === "output-available"
63+
);
64+
}
65+
66+
export function hasContextText(output: unknown): boolean {
67+
if (!isRecord(output)) return false;
68+
const contextText = output[CONTEXT_TEXT_FIELD];
69+
return typeof contextText === "string" && contextText.length > 0;
70+
}
Lines changed: 42 additions & 110 deletions
Original file line numberDiff line numberDiff line change
@@ -12,78 +12,20 @@ import {
1212
commands as templateCommands,
1313
} from "@hypr/plugin-template";
1414

15-
import type { ContextRef } from "./context/entities";
16-
import type { HyprUIMessage } from "./types";
17-
18-
function isRecord(value: unknown): value is Record<string, unknown> {
19-
return typeof value === "object" && value !== null;
20-
}
21-
22-
const MAX_TOOL_STEPS = 5;
23-
const MESSAGE_WINDOW_THRESHOLD = 20;
24-
const MESSAGE_WINDOW_SIZE = 10;
25-
26-
function isContextRef(value: unknown): value is ContextRef {
27-
return (
28-
isRecord(value) &&
29-
value.kind === "session" &&
30-
typeof value.key === "string" &&
31-
typeof value.sessionId === "string" &&
32-
(value.source === undefined ||
33-
value.source === "tool" ||
34-
value.source === "manual" ||
35-
value.source === "auto-current")
36-
);
37-
}
38-
39-
function getContextRefs(metadata: unknown): ContextRef[] {
40-
if (!isRecord(metadata) || !Array.isArray(metadata.contextRefs)) {
41-
return [];
42-
}
43-
44-
return metadata.contextRefs.filter((ref): ref is ContextRef =>
45-
isContextRef(ref),
46-
);
47-
}
48-
49-
function getSessionIdsFromSearchOutput(output: unknown): string[] {
50-
if (!isRecord(output) || !Array.isArray(output.results)) {
51-
return [];
52-
}
53-
return output.results.flatMap((item) => {
54-
if (
55-
!isRecord(item) ||
56-
(typeof item.id !== "string" && typeof item.id !== "number")
57-
) {
58-
return [];
59-
}
60-
return [String(item.id)];
61-
});
62-
}
63-
64-
type ToolOutputPart = {
65-
type: `tool-${string}`;
66-
state: "output-available";
67-
output?: unknown;
68-
[key: string]: unknown;
69-
};
70-
71-
function isToolOutputPart(value: unknown): value is ToolOutputPart {
72-
return (
73-
isRecord(value) &&
74-
typeof value.type === "string" &&
75-
value.type.startsWith("tool-") &&
76-
value.state === "output-available"
77-
);
78-
}
79-
80-
function hasContextText(output: unknown): boolean {
81-
return (
82-
isRecord(output) &&
83-
typeof output.contextText === "string" &&
84-
output.contextText.length > 0
85-
);
86-
}
15+
import type { ContextRef } from "../context/entities";
16+
import { CONTEXT_TEXT_FIELD } from "../tools";
17+
import type { HyprUIMessage } from "../types";
18+
import {
19+
getContextRefs,
20+
getSessionIdsFromSearchOutput,
21+
hasContextText,
22+
isRecord,
23+
isToolOutputPart,
24+
MAX_TOOL_STEPS,
25+
MESSAGE_WINDOW_SIZE,
26+
MESSAGE_WINDOW_THRESHOLD,
27+
type ToolOutputPart,
28+
} from "./helpers";
8729

8830
export class CustomChatTransport implements ChatTransport<HyprUIMessage> {
8931
constructor(
@@ -122,6 +64,7 @@ export class CustomChatTransport implements ChatTransport<HyprUIMessage> {
12264
return null;
12365
}
12466

67+
// Rendered by Rust-side template engine via Tauri plugin
12568
const rendered = await templateCommands.render({
12669
contextBlock: { contexts },
12770
});
@@ -130,33 +73,43 @@ export class CustomChatTransport implements ChatTransport<HyprUIMessage> {
13073
return result;
13174
}
13275

133-
private async expandSearchSessionsOutput(
134-
part: ToolOutputPart,
76+
private async hydrateSearchOutput(
77+
output: unknown,
13578
cache: Map<string, string | null>,
136-
): Promise<ToolOutputPart> {
137-
if (hasContextText(part.output)) {
138-
return part;
139-
}
140-
141-
const sessionIds = getSessionIdsFromSearchOutput(part.output);
142-
if (sessionIds.length === 0) return part;
79+
): Promise<unknown> {
80+
const sessionIds = getSessionIdsFromSearchOutput(output);
81+
if (sessionIds.length === 0) return output;
14382

14483
const refs: ContextRef[] = sessionIds.map((sessionId) => ({
145-
kind: "session",
84+
kind: "session" as const,
14685
key: `session:search:${sessionId}`,
147-
source: "tool",
86+
source: "tool" as const,
14887
sessionId,
14988
}));
15089

15190
const contextText = await this.renderContextBlock(refs, cache);
152-
if (!contextText) return part;
91+
if (!contextText) return output;
92+
93+
return {
94+
...(isRecord(output) ? output : {}),
95+
[CONTEXT_TEXT_FIELD]: contextText,
96+
};
97+
}
98+
99+
private async expandSearchSessionsOutput(
100+
part: ToolOutputPart,
101+
cache: Map<string, string | null>,
102+
): Promise<ToolOutputPart> {
103+
if (hasContextText(part.output)) {
104+
return part;
105+
}
106+
107+
const output = await this.hydrateSearchOutput(part.output, cache);
108+
if (output === part.output) return part;
153109

154110
return {
155111
...part,
156-
output: {
157-
...(isRecord(part.output) ? part.output : {}),
158-
contextText,
159-
},
112+
output,
160113
};
161114
}
162115

@@ -184,28 +137,7 @@ export class CustomChatTransport implements ChatTransport<HyprUIMessage> {
184137
if (hasContextText(output)) {
185138
return output;
186139
}
187-
188-
const sessionIds = getSessionIdsFromSearchOutput(output);
189-
if (sessionIds.length === 0) {
190-
return output;
191-
}
192-
193-
const refs: ContextRef[] = sessionIds.map((sessionId) => ({
194-
kind: "session",
195-
key: `session:search:${sessionId}`,
196-
source: "tool",
197-
sessionId,
198-
}));
199-
200-
const contextText = await this.renderContextBlock(refs, cache);
201-
if (!contextText) {
202-
return output;
203-
}
204-
205-
return {
206-
...(isRecord(output) ? output : {}),
207-
contextText,
208-
};
140+
return this.hydrateSearchOutput(output, cache);
209141
},
210142
},
211143
};

apps/desktop/src/shared/main/index.tsx

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,8 @@ import { useNewNote, useNewNoteAndListen } from "./useNewNote";
2828
import { TabContentAI, TabItemAI } from "~/ai";
2929
import { TabContentCalendar, TabItemCalendar } from "~/calendar";
3030
import { ChatFloatingButton } from "~/chat/components";
31-
import { TabContentChat } from "~/chat/tab-content";
32-
import { TabItemChat } from "~/chat/tab-item";
31+
import { TabContentChat } from "~/chat/tab/tab-content";
32+
import { TabItemChat } from "~/chat/tab/tab-item";
3333
import { TabContentChatShortcut, TabItemChatShortcut } from "~/chat_shortcuts";
3434
import { TabContentContact, TabItemContact } from "~/contacts";
3535
import { TabContentHuman, TabItemHuman } from "~/contacts/humans";

0 commit comments

Comments
 (0)