diff --git a/CHANGELOG.md b/CHANGELOG.md index a1e3b82..ebfcbf5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,14 @@ --- +## [Unreleased] + +### ✨ 新功能 + +- **Offload V2 `result_ref` 恢复接口**：新增 `POST /v2/offload/read-ref`，复用现有鉴权并将 `result_ref` 绑定到当前 session；支持全文、子串和行范围读取，服务端限制返回 token 数，本地与 COS 存储后端均可使用。 + +--- + ## [1.0.1] - 2026-07-13 ### 🐛 修复 diff --git a/sdk/python/README.md b/sdk/python/README.md index 4f54149..1dbde4f 100644 --- a/sdk/python/README.md +++ b/sdk/python/README.md @@ -80,6 +80,15 @@ compacted = client.offload_compact( ) print(compacted["messages"], compacted["report"]) +# Recover an archived tool result referenced by a compacted message +ref = client.offload_read_ref( + session_id="agent_sess_123", + result_ref="offload/agent_sess_123/refs/call_1.md", + query="relevant section", + max_tokens=800, +) +print(ref["content"], ref["truncated"]) + # Read memory pipeline artifacts (e.g. persona.md, scene_blocks/*.md) raw = client.read_file("scene_blocks/工作.md") ``` @@ -122,6 +131,7 @@ asyncio.run(main()) | L3 | `write_core()` | `POST /v2/core/write` | | Offload | `offload_ingest()` | `POST /v2/offload/ingest` | | Offload | `offload_compact()` | `POST /v2/offload/compact` | +| Offload | `offload_read_ref()` | `POST /v2/offload/read-ref` | | Offload | `offload_query_mmd()` | `POST /v2/offload/query-mmd` | ## Error Handling diff --git a/sdk/python/README_CN.md b/sdk/python/README_CN.md index 7dd0011..8e356c8 100644 --- a/sdk/python/README_CN.md +++ b/sdk/python/README_CN.md @@ -80,6 +80,15 @@ compacted = client.offload_compact( ) print(compacted["messages"], compacted["report"]) +# 按压缩消息中的 result_ref 恢复归档工具结果 +ref = client.offload_read_ref( + session_id="agent_sess_123", + result_ref="offload/agent_sess_123/refs/call_1.md", + query="相关片段", + max_tokens=800, +) +print(ref["content"], ref["truncated"]) + # 读取记忆 pipeline 产物（如 persona.md、scene_blocks/*.md） raw = client.read_file("scene_blocks/工作.md") ``` @@ -122,6 +131,7 @@ asyncio.run(main()) | L3 | `write_core()` | `POST /v2/core/write` | | Offload | `offload_ingest()` | `POST /v2/offload/ingest` | | Offload | `offload_compact()` | `POST /v2/offload/compact` | +| Offload | `offload_read_ref()` | `POST /v2/offload/read-ref` | | Offload | `offload_query_mmd()` | `POST /v2/offload/query-mmd` | ## 错误处理 diff --git a/sdk/python/tencentdb_agent_memory/client.py b/sdk/python/tencentdb_agent_memory/client.py index c4347d2..c970616 100644 --- a/sdk/python/tencentdb_agent_memory/client.py +++ b/sdk/python/tencentdb_agent_memory/client.py @@ -255,7 +255,7 @@ class MemoryClient: """``POST /core/write``""" return self._stub.post(f"{_V2}/core/write", {"content": content}) - # -- Offload (Ingest + Compact + Query-MMD) ---------------------------- + # -- Offload (Ingest + Compact + Read-Ref + Query-MMD) ----------------- def offload_ingest( self, @@ -336,6 +336,46 @@ class MemoryClient: }), ) + def offload_read_ref( + self, + session_id: str, + result_ref: str, + *, + query: Optional[str] = None, + start_line: Optional[int] = None, + end_line: Optional[int] = None, + max_tokens: Optional[int] = None, + ) -> Dict[str, Any]: + """``POST /v2/offload/read-ref`` — 读取已归档的原始工具结果。 + + 服务端会验证 ``result_ref`` 是否属于 ``session_id``，并限制返回内容长度。 + ``query`` 与行范围参数不能同时使用。 + + Parameters + ---------- + session_id : str + 归档结果所属的会话 ID。 + result_ref : str + Offload V2 压缩结果中返回的引用。 + query : str, optional + 用于定位片段的子串，匹配时不区分大小写。 + start_line, end_line : int, optional + 从 1 开始且包含首尾的行范围。 + max_tokens : int, optional + 返回内容的最大 token 数，服务端另有硬上限。 + """ + return self._stub.post( + f"{_V2}/offload/read-ref", + _strip_none({ + "session_id": session_id, + "result_ref": result_ref, + "query": query, + "start_line": start_line, + "end_line": end_line, + "max_tokens": max_tokens, + }), + ) + def offload_query_mmd( self, session_id: str, @@ -542,7 +582,7 @@ class AsyncMemoryClient: async def write_core(self, content: str) -> Dict[str, Any]: return await self._stub.post(f"{_V2}/core/write", {"content": content}) - # -- Offload (Ingest + Compact + Query-MMD) ---------------------------- + # -- Offload (Ingest + Compact + Read-Ref + Query-MMD) ----------------- async def offload_ingest( self, @@ -586,6 +626,29 @@ class AsyncMemoryClient: }), ) + async def offload_read_ref( + self, + session_id: str, + result_ref: str, + *, + query: Optional[str] = None, + start_line: Optional[int] = None, + end_line: Optional[int] = None, + max_tokens: Optional[int] = None, + ) -> Dict[str, Any]: + """``POST /v2/offload/read-ref``（异步）""" + return await self._stub.post( + f"{_V2}/offload/read-ref", + _strip_none({ + "session_id": session_id, + "result_ref": result_ref, + "query": query, + "start_line": start_line, + "end_line": end_line, + "max_tokens": max_tokens, + }), + ) + async def offload_query_mmd( self, session_id: str, diff --git a/sdk/python/tests/test_offload_read_ref.py b/sdk/python/tests/test_offload_read_ref.py new file mode 100644 index 0000000..d991f48 --- /dev/null +++ b/sdk/python/tests/test_offload_read_ref.py @@ -0,0 +1,85 @@ +from typing import Any, Dict, Optional + +import pytest + +from tencentdb_agent_memory import AsyncMemoryClient, MemoryClient +from tencentdb_agent_memory._http import Stub + + +class RecordingStub(Stub): + def __init__(self) -> None: + self.path = "" + self.body: Dict[str, Any] = {} + + def post( + self, + path: str, + body: dict, + timeout: Optional[float] = None, + ) -> dict: + self.path = path + self.body = body + return {"content": "archived result", "truncated": False} + + def close(self) -> None: + pass + + +class AsyncRecordingStub: + def __init__(self) -> None: + self.path = "" + self.body: Dict[str, Any] = {} + + async def post( + self, + path: str, + body: dict, + timeout: Optional[float] = None, + ) -> dict: + self.path = path + self.body = body + return {"content": "archived result", "truncated": False} + + +def test_offload_read_ref_posts_bounded_query_options() -> None: + stub = RecordingStub() + client = MemoryClient(stub=stub) + + result = client.offload_read_ref( + "session-1", + "offload/session-1/refs/call-1.md", + query="needle", + max_tokens=800, + ) + + assert stub.path == "/v2/offload/read-ref" + assert stub.body == { + "session_id": "session-1", + "result_ref": "offload/session-1/refs/call-1.md", + "query": "needle", + "max_tokens": 800, + } + assert result["content"] == "archived result" + + +@pytest.mark.asyncio +async def test_async_offload_read_ref_posts_line_range() -> None: + stub = AsyncRecordingStub() + client = AsyncMemoryClient.__new__(AsyncMemoryClient) + client._stub = stub + + result = await client.offload_read_ref( + "session-1", + "offload/session-1/refs/call-1.md", + start_line=5, + end_line=12, + ) + + assert stub.path == "/v2/offload/read-ref" + assert stub.body == { + "session_id": "session-1", + "result_ref": "offload/session-1/refs/call-1.md", + "start_line": 5, + "end_line": 12, + } + assert result["content"] == "archived result" diff --git a/sdk/typescript/README.md b/sdk/typescript/README.md index 73ce541..859d33f 100644 --- a/sdk/typescript/README.md +++ b/sdk/typescript/README.md @@ -76,6 +76,15 @@ const compacted = await client.offloadCompact({ }); console.log(compacted.messages, compacted.report); +// Recover an archived tool result referenced by a compacted message +const ref = await client.offloadReadRef({ + session_id: "agent_sess_123", + result_ref: "offload/agent_sess_123/refs/call_1.md", + query: "relevant section", + max_tokens: 800, +}); +console.log(ref.content, ref.truncated); + // Offload v2: query MMD task graphs const mmd = await client.offloadQueryMmd({ session_id: "agent_sess_123", limit: 1 }); console.log(mmd.current_mmd, mmd.mmds); @@ -104,6 +113,7 @@ const raw = await client.readFile("scene_blocks/工作.md"); | L3 | `writeCore()` | `POST /v2/core/write` | | Offload | `offloadIngest()` | `POST /v2/offload/ingest` | | Offload | `offloadCompact()` | `POST /v2/offload/compact` | +| Offload | `offloadReadRef()` | `POST /v2/offload/read-ref` | | Offload | `offloadQueryMmd()` | `POST /v2/offload/query-mmd` | ## Error Handling diff --git a/sdk/typescript/README_CN.md b/sdk/typescript/README_CN.md index 78a19a1..6f7fb65 100644 --- a/sdk/typescript/README_CN.md +++ b/sdk/typescript/README_CN.md @@ -76,6 +76,15 @@ const compacted = await client.offloadCompact({ }); console.log(compacted.messages, compacted.report); +// 按压缩消息中的 result_ref 恢复归档工具结果 +const ref = await client.offloadReadRef({ + session_id: "agent_sess_123", + result_ref: "offload/agent_sess_123/refs/call_1.md", + query: "相关片段", + max_tokens: 800, +}); +console.log(ref.content, ref.truncated); + // Offload v2: 查询任务流程图（MMD） const mmd = await client.offloadQueryMmd({ session_id: "agent_sess_123", limit: 1 }); console.log(mmd.current_mmd, mmd.mmds); @@ -104,6 +113,7 @@ const raw = await client.readFile("scene_blocks/工作.md"); | L3 | `writeCore()` | `POST /v2/core/write` | | Offload | `offloadIngest()` | `POST /v2/offload/ingest` | | Offload | `offloadCompact()` | `POST /v2/offload/compact` | +| Offload | `offloadReadRef()` | `POST /v2/offload/read-ref` | | Offload | `offloadQueryMmd()` | `POST /v2/offload/query-mmd` | ## 错误处理 diff --git a/sdk/typescript/src/client.ts b/sdk/typescript/src/client.ts index 16ed571..98db4f0 100644 --- a/sdk/typescript/src/client.ts +++ b/sdk/typescript/src/client.ts @@ -1,7 +1,7 @@ /** * TencentDB Agent Memory v2 TypeScript SDK — `MemoryClient`. * - * 14 methods mapping 1:1 to the v2 data-plane API. + * 18 methods mapping 1:1 to the v2 data-plane API. */ import { HttpTransport, type HttpTransportOptions } from "./http.js"; @@ -29,6 +29,8 @@ import type { OffloadCompactRequest, OffloadIngestData, OffloadIngestRequest, + OffloadReadRefData, + OffloadReadRefRequest, OffloadQueryMmdData, OffloadQueryMmdRequest, ScenarioFile, @@ -154,7 +156,7 @@ export class MemoryClient { return this.http.post(`${V2}/core/write`, params as unknown as Record); } - // -- Offload (Compaction + Ingest) ------------------------------------ + // -- Offload ----------------------------------------------------------- /** * Send tool pairs (+ optional context) to offload server for L1 processing. @@ -172,6 +174,13 @@ export class MemoryClient { return this.http.post(`${V2}/offload/compact`, stripUndefined(params as unknown as Record)); } + /** + * Recover an archived tool result referenced by a compacted message. + */ + offloadReadRef(params: OffloadReadRefRequest): Promise { + return this.http.post(`${V2}/offload/read-ref`, stripUndefined(params as unknown as Record)); + } + /** * Query MMD task graphs for a session. * limit=1 returns only the current active MMD (fast path). diff --git a/sdk/typescript/src/index.ts b/sdk/typescript/src/index.ts index d1b319a..dc54563 100644 --- a/sdk/typescript/src/index.ts +++ b/sdk/typescript/src/index.ts @@ -28,6 +28,8 @@ export type { OffloadToolPair, OffloadRecentMessage, OffloadIngestRequest, OffloadIngestData, OffloadCompactRequest, OffloadCompactData, OffloadCompactReport, + OffloadReadRefRequest, OffloadReadRefData, + OffloadQueryMmdRequest, OffloadQueryMmdData, // Common ApiResponseEnvelope, } from "./types.js"; diff --git a/sdk/typescript/src/types.ts b/sdk/typescript/src/types.ts index be0f435..9645295 100644 --- a/sdk/typescript/src/types.ts +++ b/sdk/typescript/src/types.ts @@ -245,6 +245,22 @@ export interface OffloadCompactData { report: OffloadCompactReport; } +export interface OffloadReadRefRequest { + session_id: string; + result_ref: string; + query?: string; + start_line?: number; + end_line?: number; + max_tokens?: number; +} + +export interface OffloadReadRefData { + result_ref: string; + content: string; + truncated: boolean; + match_found?: boolean; +} + export interface OffloadQueryMmdRequest { session_id: string; limit?: number; diff --git a/sdk/typescript/tests/client.test.ts b/sdk/typescript/tests/client.test.ts index 4686eae..e2d95fb 100644 --- a/sdk/typescript/tests/client.test.ts +++ b/sdk/typescript/tests/client.test.ts @@ -37,6 +37,11 @@ function createMock(): MockTransport { "/v2/scenario/rm": {}, "/v2/core/read": { content: "# core", created_at: "t", updated_at: "t" }, "/v2/core/write": { updated_at: "t" }, + "/v2/offload/read-ref": { + result_ref: "offload/s1/refs/call-1.md", + content: "archived result", + truncated: false, + }, }; return m; } @@ -187,6 +192,32 @@ describe("L3 Core", () => { }); }); +// --------------------------------------------------------------------------- +// Offload +// --------------------------------------------------------------------------- + +describe("Offload", () => { + it("offloadReadRef sends the reference and strips undefined filters", async () => { + const mock = createMock(); + const client = new MemoryClient(mock); + const result = await client.offloadReadRef({ + session_id: "s1", + result_ref: "offload/s1/refs/call-1.md", + query: "result", + max_tokens: 800, + }); + + expect(result.content).toBe("archived result"); + expect(mock.calls[0]!.path).toBe("/v2/offload/read-ref"); + expect(mock.calls[0]!.body).toEqual({ + session_id: "s1", + result_ref: "offload/s1/refs/call-1.md", + query: "result", + max_tokens: 800, + }); + }); +}); + // --------------------------------------------------------------------------- // Error handling // --------------------------------------------------------------------------- diff --git a/src/gateway/server.ts b/src/gateway/server.ts index 4616046..f8a9ca1 100644 --- a/src/gateway/server.ts +++ b/src/gateway/server.ts @@ -601,7 +601,7 @@ export class TdaiGateway { } } - // ── Offload V2 routes (async ingest + mmd query) ── + // ── Offload V2 routes (ingest, compact, ref recovery, MMD query) ── const offloadDeps: OffloadV2Deps = { resolveStorage: v2Deps.resolveStorage, getStorage: v2Deps.getStorage ?? (() => undefined), diff --git a/src/offload_server/compact/helpers.ts b/src/offload_server/compact/helpers.ts index d9dc902..39afb8d 100644 --- a/src/offload_server/compact/helpers.ts +++ b/src/offload_server/compact/helpers.ts @@ -184,7 +184,10 @@ export function replaceWithSummary(msg: Message, entry: OffloadEntry): void { `Summary: ${entry.summary}`, ]; if (entry.result_ref) { - parts.push(`原始工具结果已存档，如需查看完整内容请调用 tdai_read_cos(path="${entry.result_ref}")`); + parts.push( + "原始工具结果已存档，如需查看完整内容请调用 Offload V2 " + + `result_ref 恢复接口（POST /v2/offload/read-ref，result_ref="${entry.result_ref}"）`, + ); } const summaryContent = parts.join("\n"); diff --git a/src/offload_server/offload-task-executor.ts b/src/offload_server/offload-task-executor.ts index cd51b9d..fb25293 100644 --- a/src/offload_server/offload-task-executor.ts +++ b/src/offload_server/offload-task-executor.ts @@ -213,7 +213,7 @@ export class OffloadTaskExecutor { const header = `# Tool Result: ${toolName}\n\n**tool_call_id:** ${id}\n**Timestamp:** ${timestamp}\n\n---\n\n`; const refPath = `${basePath}/refs/${id}.md`; await storage.writeFile(refPath, header + resultStr); - // Set result_ref on the matching entry (full relative path for tdai_read_cos) + // Set result_ref on the matching entry for bounded recovery through Offload V2. const entry = newEntries.find((e) => e.tool_call_id === id); if (entry) entry.result_ref = `${basePath}/refs/${id}.md`; } diff --git a/src/offload_server/read-ref-handler.test.ts b/src/offload_server/read-ref-handler.test.ts new file mode 100644 index 0000000..fed8f29 --- /dev/null +++ b/src/offload_server/read-ref-handler.test.ts @@ -0,0 +1,260 @@ +import type http from "node:http"; +import { getEncoding } from "js-tiktoken"; +import { describe, expect, it, vi } from "vitest"; +import type { StorageAdapter } from "../core/storage/adapter.js"; +import { replaceWithSummary } from "./compact/helpers.js"; +import { handleOffloadV2Route } from "./router.js"; +import { + handleReadRef, + resolveOwnedResultRef, + sliceRefContent, +} from "./read-ref-handler.js"; + +const requestId = "req-test"; +const auth = { serviceId: "service-test" }; + +function envelopes() { + return { + successEnvelope: (data: T, id: string) => ({ code: 0, message: "ok", request_id: id, data }), + errorEnvelope: (code: number, message: string, id: string) => ({ + code, + message, + request_id: id, + }), + }; +} + +function storageWith(content: string | null) { + const readFile = vi.fn(async () => content); + return { + storage: { readFile } as unknown as StorageAdapter, + readFile, + }; +} + +describe("resolveOwnedResultRef", () => { + it("accepts only direct refs belonging to the requested session", () => { + expect(resolveOwnedResultRef("session-1", "offload/session-1/refs/call-1.md")) + .toBe("offload/session-1/refs/call-1.md"); + expect(resolveOwnedResultRef("session-1", "offload/session-2/refs/call-1.md")).toBeNull(); + expect(resolveOwnedResultRef("session-1", "offload/session-1/refs/../secret.md")).toBeNull(); + expect(resolveOwnedResultRef("session-1", "/offload/session-1/refs/call-1.md")).toBeNull(); + expect(resolveOwnedResultRef("session-1", "offload/session-1/refs/nested/call-1.md")).toBeNull(); + expect(resolveOwnedResultRef("session-1", "offload/session-1/refs/call-1.txt")).toBeNull(); + }); + + it("uses the same sanitized session path as offload writes", () => { + expect(resolveOwnedResultRef("agent:session", "offload/agent_session/refs/call-1.md")) + .toBe("offload/agent_session/refs/call-1.md"); + }); +}); + +describe("sliceRefContent", () => { + it("reads an inclusive line range and reports partial content", () => { + expect(sliceRefContent("one\ntwo\nthree\nfour", { + start_line: 2, + end_line: 3, + max_tokens: 100, + })).toEqual({ + content: "two\nthree", + truncated: true, + }); + }); + + it("returns a bounded excerpt around a query", () => { + const raw = `${"before ".repeat(200)}TARGET${" after".repeat(200)}`; + const result = sliceRefContent(raw, { + query: "target", + max_tokens: 32, + }); + + expect(result.content).toContain("TARGET"); + expect(result.truncated).toBe(true); + expect(result.match_found).toBe(true); + expect(getEncoding("o200k_base").encode(result.content).length).toBeLessThanOrEqual(32); + }); + + it("bounds an unfiltered result by the requested token budget", () => { + const result = sliceRefContent("token ".repeat(100), { + max_tokens: 5, + }); + + expect(result.truncated).toBe(true); + expect(getEncoding("o200k_base").encode(result.content).length).toBeLessThanOrEqual(5); + }); + + it("bounds a query that is itself larger than the token budget", () => { + const query = "needle ".repeat(20); + const result = sliceRefContent(`before ${query} after`, { + query, + max_tokens: 3, + }); + + expect(result.truncated).toBe(true); + expect(result.match_found).toBe(true); + expect(getEncoding("o200k_base").encode(result.content).length).toBeLessThanOrEqual(3); + }); + + it("returns an explicit no-match result", () => { + expect(sliceRefContent("alpha\nbeta", { + query: "missing", + max_tokens: 100, + })).toEqual({ + content: "", + truncated: false, + match_found: false, + }); + }); +}); + +describe("compaction recovery hint", () => { + it("keeps the existing Chinese wording while pointing to the V2 route", () => { + const message: any = { role: "tool", content: "raw result" }; + + replaceWithSummary(message, { + tool_call_id: "call-1", + tool_call: "search", + summary: "result summary", + timestamp: "2026-07-24T00:00:00Z", + score: 2, + node_id: "node-1", + result_ref: "offload/session-1/refs/call-1.md", + }); + + expect(message.content).toContain("原始工具结果已存档，如需查看完整内容请调用"); + expect(message.content).toContain("POST /v2/offload/read-ref"); + expect(message.content).not.toContain("tdai_read_cos"); + }); +}); + +describe("handleReadRef", () => { + it("reads an owned reference and wraps the response", async () => { + const { storage, readFile } = storageWith("archived result"); + let sent: { status: number; body: any } | undefined; + const { successEnvelope, errorEnvelope } = envelopes(); + + await handleReadRef( + {} as http.IncomingMessage, + {} as http.ServerResponse, + auth, + storage, + requestId, + async () => ({ + session_id: "session-1", + result_ref: "offload/session-1/refs/call-1.md", + }), + (_res, status, body) => { sent = { status, body }; }, + successEnvelope, + errorEnvelope, + ); + + expect(readFile).toHaveBeenCalledWith("offload/session-1/refs/call-1.md"); + expect(sent).toEqual({ + status: 200, + body: { + code: 0, + message: "ok", + request_id: requestId, + data: { + result_ref: "offload/session-1/refs/call-1.md", + content: "archived result", + truncated: false, + }, + }, + }); + }); + + it("hides invalid, cross-session, and missing references behind 404", async () => { + const { successEnvelope, errorEnvelope } = envelopes(); + for (const testCase of [ + { + resultRef: "offload/other-session/refs/call-1.md", + stored: "must not be read", + expectedReads: 0, + }, + { + resultRef: "offload/session-1/refs/missing.md", + stored: null, + expectedReads: 1, + }, + ]) { + const { storage, readFile } = storageWith(testCase.stored); + let sent: { status: number; body: any } | undefined; + + await handleReadRef( + {} as http.IncomingMessage, + {} as http.ServerResponse, + auth, + storage, + requestId, + async () => ({ + session_id: "session-1", + result_ref: testCase.resultRef, + }), + (_res, status, body) => { sent = { status, body }; }, + successEnvelope, + errorEnvelope, + ); + + expect(readFile).toHaveBeenCalledTimes(testCase.expectedReads); + expect(sent?.status).toBe(404); + expect(sent?.body.message).toBe("result_ref not found"); + } + }); + + it("rejects requests above the server token cap", async () => { + const { storage, readFile } = storageWith("must not be read"); + let sent: { status: number; body: any } | undefined; + const { successEnvelope, errorEnvelope } = envelopes(); + + await handleReadRef( + {} as http.IncomingMessage, + {} as http.ServerResponse, + auth, + storage, + requestId, + async () => ({ + session_id: "session-1", + result_ref: "offload/session-1/refs/call-1.md", + max_tokens: 4097, + }), + (_res, status, body) => { sent = { status, body }; }, + successEnvelope, + errorEnvelope, + ); + + expect(readFile).not.toHaveBeenCalled(); + expect(sent?.status).toBe(400); + }); + + it("is dispatched by the authenticated Offload V2 router", async () => { + const { storage } = storageWith("routed content"); + const req = { + headers: { + authorization: "Bearer test-key", + "x-tdai-service-id": "service-test", + }, + } as http.IncomingMessage; + let sent: { status: number; body: any } | undefined; + + const handled = await handleOffloadV2Route( + req, + {} as http.ServerResponse, + "/v2/offload/read-ref/", + "POST", + async () => ({ + session_id: "session-1", + result_ref: "offload/session-1/refs/call-1.md", + }), + (_res, status, body) => { sent = { status, body }; }, + { + getStorage: () => storage, + logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn() }, + }, + ); + + expect(handled).toBe(true); + expect(sent?.status).toBe(200); + expect(sent?.body.data.content).toBe("routed content"); + }); +}); diff --git a/src/offload_server/read-ref-handler.ts b/src/offload_server/read-ref-handler.ts new file mode 100644 index 0000000..6ea8e8d --- /dev/null +++ b/src/offload_server/read-ref-handler.ts @@ -0,0 +1,199 @@ +/** + * Offload Result Reference Handler — bounded recovery of archived tool results. + */ +import type http from "node:http"; +import { getEncoding, type Tiktoken } from "js-tiktoken"; +import type { StorageAdapter } from "../core/storage/adapter.js"; +import { ReadRefRequestSchema, type ReadRefRequest } from "./schemas.js"; +import { buildOffloadBasePath } from "./session-utils.js"; + +let encoder: Tiktoken | null = null; + +function getEncoder(): Tiktoken { + if (!encoder) encoder = getEncoding("o200k_base"); + return encoder; +} + +function countTokens(text: string): number { + if (!text) return 0; + try { + return getEncoder().encode(text).length; + } catch { + // Byte length is a conservative upper bound for byte-level BPE tokens. + return Buffer.byteLength(text, "utf-8"); + } +} + +function truncatePrefix(text: string, maxTokens: number): { content: string; truncated: boolean } { + if (countTokens(text) <= maxTokens) { + return { content: text, truncated: false }; + } + + const characters = Array.from(text); + let low = 0; + let high = characters.length; + while (low < high) { + const mid = Math.ceil((low + high) / 2); + if (countTokens(characters.slice(0, mid).join("")) <= maxTokens) { + low = mid; + } else { + high = mid - 1; + } + } + + return { + content: characters.slice(0, low).join(""), + truncated: true, + }; +} + +function sliceAroundQuery( + raw: string, + query: string, + maxTokens: number, +): { content: string; truncated: boolean; matchFound: boolean } { + const matchIndex = raw.toLowerCase().indexOf(query.toLowerCase()); + if (matchIndex < 0) { + return { content: "", truncated: false, matchFound: false }; + } + + const characters = Array.from(raw); + const matchStart = Array.from(raw.slice(0, matchIndex)).length; + const matchEnd = matchStart + Array.from(raw.slice(matchIndex, matchIndex + query.length)).length; + const matchText = characters.slice(matchStart, matchEnd).join(""); + + if (countTokens(matchText) > maxTokens) { + const limited = truncatePrefix(matchText, maxTokens); + return { ...limited, matchFound: true }; + } + + let low = 0; + let high = Math.max(matchStart, characters.length - matchEnd); + let bestStart = matchStart; + let bestEnd = matchEnd; + + while (low <= high) { + const radius = Math.floor((low + high) / 2); + const start = Math.max(0, matchStart - radius); + const end = Math.min(characters.length, matchEnd + radius); + const candidate = characters.slice(start, end).join(""); + if (countTokens(candidate) <= maxTokens) { + bestStart = start; + bestEnd = end; + low = radius + 1; + } else { + high = radius - 1; + } + } + + return { + content: characters.slice(bestStart, bestEnd).join(""), + truncated: bestStart > 0 || bestEnd < characters.length, + matchFound: true, + }; +} + +export interface ReadRefData { + result_ref: string; + content: string; + truncated: boolean; + match_found?: boolean; +} + +/** + * Resolve a result reference only when it points to a direct Markdown child + * of the current session's refs directory. + */ +export function resolveOwnedResultRef(sessionId: string, resultRef: string): string | null { + const prefix = `${buildOffloadBasePath(sessionId)}/refs/`; + if (!resultRef.startsWith(prefix)) return null; + + const filename = resultRef.slice(prefix.length); + if ( + !filename || + filename.includes("/") || + filename.includes("\\") || + filename.includes("\0") || + filename === "." || + filename === ".." || + !filename.endsWith(".md") + ) { + return null; + } + return resultRef; +} + +/** + * Select and bound content from an archived tool-result reference. + */ +export function sliceRefContent( + raw: string, + request: Pick, +): Omit { + if (request.query) { + const result = sliceAroundQuery(raw, request.query, request.max_tokens); + return { + content: result.content, + truncated: result.truncated, + match_found: result.matchFound, + }; + } + + const lines = raw.split(/\r?\n/u); + const startIndex = Math.min((request.start_line ?? 1) - 1, lines.length); + const endIndex = Math.min(request.end_line ?? lines.length, lines.length); + const selected = lines.slice(startIndex, endIndex).join("\n"); + const limited = truncatePrefix(selected, request.max_tokens); + + return { + content: limited.content, + truncated: limited.truncated || startIndex > 0 || endIndex < lines.length, + }; +} + +/** + * Handle POST /v2/offload/read-ref. + */ +export async function handleReadRef( + req: http.IncomingMessage, + res: http.ServerResponse, + _auth: { serviceId: string }, + storage: StorageAdapter, + requestId: string, + parseJsonBody: (req: http.IncomingMessage) => Promise, + sendJson: (res: http.ServerResponse, status: number, body: unknown) => void, + successEnvelope: (data: T, requestId: string) => unknown, + errorEnvelope: (code: number, message: string, requestId: string) => unknown, +): Promise { + const body = await parseJsonBody(req); + const parsed = ReadRefRequestSchema.safeParse(body); + if (!parsed.success) { + sendJson(res, 400, errorEnvelope(400, parsed.error.message, requestId)); + return; + } + + const { session_id: sessionId, result_ref: resultRef } = parsed.data; + const ownedRef = resolveOwnedResultRef(sessionId, resultRef); + if (!ownedRef) { + sendJson(res, 404, errorEnvelope(404, "result_ref not found", requestId)); + return; + } + + const raw = await storage.readFile(ownedRef); + if (raw === null) { + sendJson(res, 404, errorEnvelope(404, "result_ref not found", requestId)); + return; + } + + sendJson( + res, + 200, + successEnvelope( + { + result_ref: ownedRef, + ...sliceRefContent(raw, parsed.data), + }, + requestId, + ), + ); +} diff --git a/src/offload_server/router.ts b/src/offload_server/router.ts index 783916f..052e18e 100644 --- a/src/offload_server/router.ts +++ b/src/offload_server/router.ts @@ -10,6 +10,7 @@ import { parseV2Auth, successEnvelope, errorEnvelope, makeRequestId } from "../g import { handleIngest } from "./ingest-handler.js"; import { handleMmdQuery } from "./mmd-handler.js"; import { handleCompaction } from "./compact/compaction-handler.js"; +import { handleReadRef } from "./read-ref-handler.js"; import { MmdQuerySchema } from "./schemas.js"; export interface OffloadV2Deps { @@ -82,6 +83,20 @@ export async function handleOffloadV2Route( }, requestId, parseJsonBody, sendJson, successEnvelope, errorEnvelope); return true; + case "POST /v2/offload/read-ref": + await handleReadRef( + req, + res, + auth, + storage, + requestId, + parseJsonBody, + sendJson, + successEnvelope, + errorEnvelope, + ); + return true; + default: return false; } diff --git a/src/offload_server/schemas.ts b/src/offload_server/schemas.ts index 634b38e..a559d57 100644 --- a/src/offload_server/schemas.ts +++ b/src/offload_server/schemas.ts @@ -78,3 +78,23 @@ export const MmdQuerySchema = z.object({ session_id: safeSessionId, limit: z.number().int().min(1).optional(), }); + +export const ReadRefRequestSchema = z + .object({ + session_id: safeSessionId, + result_ref: z.string().trim().min(1).max(1000), + query: z.string().trim().min(1).max(1000).optional(), + start_line: z.number().int().min(1).optional(), + end_line: z.number().int().min(1).optional(), + max_tokens: z.number().int().min(1).max(4096).default(1600), + }) + .refine( + (data) => data.start_line === undefined || data.end_line === undefined || data.start_line <= data.end_line, + { message: "start_line must not exceed end_line" }, + ) + .refine( + (data) => data.query === undefined || (data.start_line === undefined && data.end_line === undefined), + { message: "query cannot be combined with start_line or end_line" }, + ); + +export type ReadRefRequest = z.infer; diff --git a/src/offload_server/types.ts b/src/offload_server/types.ts index 9d375a4..223f61c 100644 --- a/src/offload_server/types.ts +++ b/src/offload_server/types.ts @@ -24,7 +24,7 @@ export interface OffloadEntry { timestamp: string; score: number; node_id: string | null; - /** Full relative path to ref file storing original tool result (e.g. "offload/{sessionId}/refs/call_214.md"), readable via tdai_read_cos */ + /** Full relative path to an archived tool result, readable through POST /v2/offload/read-ref. */ result_ref?: string; }