8000 feat(history): get_messages 可按 ref 分页读回当前会话的原始工具结果 (#1650) · felinics/Memoh@79074b7 · GitHub
[go: up one dir, main page]

Skip to content

Commit 79074b7

Browse files
authored
feat(history): get_messages 可按 ref 分页读回当前会话的原始工具结果 (#1650)
1 parent cf4bcf0 commit 79074b7

8 files changed

Lines changed: 444 additions & 23 deletions

File tree

‎.github/workflows/go-ci.yml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,8 @@ jobs:
218218
-run '^TestPostgresDiscuss(RecoveryPreservesBatchAcrossCompactorRestart|BacklogOverBudgetIsAnsweredAndNextRoundContinues|BacklogWithoutCompactionIsAnsweredAndRecordsOmission|OversizedInputFailsOnceThenNextMessageIsAnswered|ImageBacklogKeepsImagesOnTheirMessages|OversizedImageInputFailsOnceThenNextMessageIsAnswered|ImageBacklogYieldsAnsweredHistoryFirst|CompactedImageInputIsRecorded|CompactionKeepsBatchTextOverImages)$' ./cmd/agent
219219
TEST_POSTGRES_BOOTSTRAP_SCHEMA=0 go test -v -race -count=1 -timeout 90s \
220220
-run '^TestPostgres(SummaryOnlyRollup|CompleteCompactionRollup)' ./internal/chat/message
221+
TEST_POSTGRES_BOOTSTRAP_SCHEMA=0 go test -v -race -count=1 -timeout 90s \
222+
-run '^TestPostgresHistoryToolResultRef$' ./internal/agent/tool
221223
- name: Verify atomic connector disconnect and tenant isolation
222224
env:
223225
TEST_POSTGRES_DSN: postgres://postgres:postgres@127.0.0.1:${{ job.services.postgres.ports['5432'] }}/memoh_runtime_test?sslmode=disable

‎cmd/internal/core/providers.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1191,7 +1191,7 @@ func provideToolProviders(log *slog.Logger, channelRuntime channel.Runtime, regi
11911191
agenttools.NewVideoGenProvider(log, settingsService, videoService, bgManager, manager, config.DefaultDataMount),
11921192
agenttools.NewFederationProvider(log, connectorSource),
11931193
agenttools.NewFederationProvider(log, fedSource),
1194-
agenttools.NewHistoryProvider(log, historySessions, messageService, queries),
1194+
agenttools.NewHistoryProvider(log, historySessions, messageService, queries, agentLimitsFromConfig(cfg.Agent).ToolOutputLimit()),
11951195
}
11961196
}
11971197

‎internal/agent/tool/history.go‎

Lines changed: 153 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,20 @@ import (
44
"context"
55
"encoding/json"
66
"errors"
7+
"fmt"
78
"log/slog"
9+
"sort"
10+
"strconv"
811
"strings"
912
"time"
13+
"unicode/utf8"
1014

1115
sdk "github.com/felinics/twilight/sdk"
1216
"github.com/jackc/pgx/v5"
1317
"github.com/jackc/pgx/v5/pgtype"
1418

19+
historyfrag "github.com/felinics/memoh/internal/agent/context/history"
20+
contextlimit "github.com/felinics/memoh/internal/agent/context/limit"
1521
"github.com/felinics/memoh/internal/agent/toolexec"
1622
"github.com/felinics/memoh/internal/agent/turn"
1723
messagepkg "github.com/felinics/memoh/internal/chat/message"
@@ -37,21 +43,25 @@ type HistoryMessageReader interface {
3743

3844
// HistoryProvider exposes list_sessions, get_messages, and search_messages tools.
3945
type HistoryProvider struct {
40-
sessions SessionLister
41-
messages HistoryMessageReader
42-
queries dbstore.Queries
43-
logger *slog.Logger
46+
sessions SessionLister
47+
messages HistoryMessageReader
48+
queries dbstore.Queries
49+
outputLimit ToolOutputLimit
50+
logger *slog.Logger
4451
}
4552

46-
func NewHistoryProvider(log *slog.Logger, sessions SessionLister, messages HistoryMessageReader, queries dbstore.Queries) *HistoryProvider {
53+
// NewHistoryProvider takes the tool output limit the runtime wraps every tool
54+
// in, so a page of a stored tool result is cut to fit it rather than elided.
55+
func NewHistoryProvider(log *slog.Logger, sessions SessionLister, messages HistoryMessageReader, queries dbstore.Queries, outputLimit ToolOutputLimit) *HistoryProvider {
4756
if log == nil {
4857
log = slog.Default()
4958
}
5059
return &HistoryProvider{
51-
sessions: sessions,
52-
messages: messages,
53-
queries: queries,
54-
logger: log.With(slog.String("tool", "history")),
60+
sessions: sessions,
61+
messages: messages,
62+
queries: queries,
63+
outputLimit: outputLimit,
64+
logger: log.With(slog.String("tool", "history")),
5565
}
5666
}
5767

@@ -63,7 +73,7 @@ func (*HistoryProvider) Usage(_ context.Context, _ SessionContext, available Ava
6373
parts = append(parts, ref+": List accessible chat sessions with their bound contact/route info. Filter by `type` (chat/schedule) or `platform`.")
6474
}
6575
if ref, ok := available.Ref(ToolGetMessages()); ok {
66-
parts = append(parts, ref+": Get recent messages from the current or selected session, or resolve one exact `message_id`.")
76+
parts = append(parts, ref+": Get recent messages from the current or selected session, resolve one exact `message_id`, or page through an earlier tool result a summary cites by `ref`.")
6777
if listSessionsRef != "" {
6878
parts = append(parts, "Use session IDs from "+listSessionsRef+" as `session_id` for "+ref+" when reading a specific conversation.")
6979
}
@@ -96,8 +106,8 @@ func (p *HistoryProvider) Tools(_ context.Context, sess SessionContext) ([]toole
96106
s := sess
97107
tools = append(tools, toolexec.Tool{
98108
Name: ToolGetMessages().String(),
99-
Description: "Get recent messages from a chat session, or resolve one exact message ID. Defaults to the current session. Results are returned oldest-first.",
100-
Parameters: toolexec.SchemaFor[getMessagesArgs](),
109+
Description: "Get recent messages from a chat session, or resolve one exact message ID. Defaults to the current session. Results are returned oldest-first. When a summary cites an earlier tool result as ref=<message id>:<part index> and you need its original output, pass that ref: the stored result text of the current session comes back in pages with total_bytes; call again with offset set to next_offset until next_offset is absent.",
110+
Parameters: toolexec.SchemaFor[getMessagesArgs](toolexec.Minimum("offset", 0), toolexec.Minimum("max_bytes", 1)),
101111
Execute: toolexec.Typed(func(ctx *toolexec.ToolExecContext, args getMessagesArgs) (sdk.ToolOutput, error) {
102112
return toolexec.OutputPair(p.execGetMessages(ctx.Context, s, args))
103113
}),
@@ -197,6 +207,12 @@ func (p *HistoryProvider) execGetMessages(ctx context.Context, sess SessionConte
197207
if botID == "" {
198208
return nil, errors.New("bot_id is required")
199209
}
210+
if strings.TrimSpace(args.Ref) != "" {
211+
return p.execGetToolResult(ctx, sess, args)
212+
}
213+
if args.Offset != nil || args.MaxBytes != nil {
214+
return nil, errors.New("offset and max_bytes page a tool result and require ref")
215+
}
200216

201217
limit := int32(30)
202218
if args.Limit != nil && *args.Limit > 0 {
@@ -271,6 +287,128 @@ func (p *HistoryProvider) execGetMessages(ctx context.Context, sess SessionConte
271287
return out, nil
272288
}
273289

290+
// execGetToolResult pages through the text a provider receives when history
291+
// replays one stored tool result: the text output itself, or its canonical
292+
// JSON document. A ref names the row and the part's position in the replayed
293+
// message, and resolves only inside the current session.
294+
func (p *HistoryProvider) execGetToolResult(ctx context.Context, sess SessionContext, args getMessagesArgs) (any, error) {
295+
if strings.TrimSpace(args.MessageID) != "" || strings.TrimSpace(args.Before) != "" || args.Limit != nil {
296+
return nil, errors.New("ref cannot be combined with message_id, before, or limit; page with offset and max_bytes")
297+
}
298+
sessionID := strings.TrimSpace(sess.SessionID)
299+
if sessionID == "" {
300+
return nil, errors.New("ref reads the current session, and there is none")
301+
}
302+
if requested := strings.TrimSpace(args.SessionID); requested != "" && requested != sessionID {
303+
return nil, errors.New("ref reads only the current session; omit session_id")
304+
}
305+
messageID, partIndex, err := parseToolResultRef(args.Ref)
306+
if err != nil {
307+
return nil, err
308+
}
309+
offset := 0
310+
if args.Offset != nil {
311+
offset = *args.Offset
312+
}
313+
if offset < 0 {
314+
return nil, errors.New("offset must be >= 0")
315+
}
316+
if args.MaxBytes != nil && *args.MaxBytes < 1 {
317+
return nil, errors.New("max_bytes must be >= 1")
318+
}
319+
320+
msg, err := p.messages.GetByIDBySession(ctx, sessionID, messageID)
321+
if errors.Is(err, pgx.ErrNoRows) {
322+
return nil, fmt.Errorf("message %s is not in the current session; ref reads only the current session", messageID)
323+
}
324+
if err != nil {
325+
return nil, err
326+
}
327+
replayed := historyfrag.StoredModelMessageToSDKMessage(historyfrag.DecodeStoredModelMessage(p.logger, msg.ID, msg.Role, msg.Content))
328+
if partIndex >= len(replayed.Content) {
329+
return nil, fmt.Errorf("message %s has no part %d", messageID, partIndex)
330+
}
331+
result, ok := replayed.Content[partIndex].(sdk.ToolResultPart)
332+
if !ok {
333+
return nil, fmt.Errorf("part %d of message %s is a %s part, not a tool result", partIndex, messageID, replayed.Content[partIndex].PartType())
334+
}
335+
text := result.Result.String()
336+
if offset > len(text) {
337+
return nil, fmt.Errorf("offset %d is past the end of the result (total_bytes %d)", offset, len(text))
338+
}
339+
return p.toolResultPage(msg.ID+":"+strconv.Itoa(partIndex), result, text, offset, args.MaxBytes), nil
340+
}
341+
342+
func parseToolResultRef(raw string) (string, int, error) {
343+
messageID, part, found := strings.Cut(strings.TrimPrefix(strings.TrimSpace(raw), "ref="), ":")
344+
index, indexErr := strconv.ParseUint(part, 10, 31)
345+
if _, idErr := dbpkg.ParseUUID(messageID); !found || indexErr != nil || idErr != nil {
346+
return "", 0, fmt.Errorf("ref must be <message id>:<part index> as a summary cites it, got %q", raw)
347+
}
348+
return messageID, int(index), nil
349+
}
350+
351+
// toolResultPage returns the longest page from offset that the runtime's
352+
// output limit leaves unchanged, so its content is exactly
353+
// text[offset:next_offset]. Pages begin and end on rune boundaries, and while
354+
// text remains a page holds at least one rune, even past max_bytes.
355+
func (p *HistoryProvider) toolResultPage(ref string, result sdk.ToolResultPart, text string, offset int, maxBytes *int) map[string]any {
356+
limit := contextlimit.NormalizedLimit(p.outputLimit)
357+
start := offset
358+
for start > 0 && start < len(text) && !utf8.RuneStart(text[start]) {
359+
start--
360+
}
361+
budget := limit.MaxBytes
362+
if maxBytes != nil {
363+
budget = min(budget, *maxBytes)
364+
}
365+
end := len(text)
366+
if budget < end-start {
367+
end = start + budget
368+
}
369+
newlines := 0
370+
for i := start; i < end; i++ {
371+
if text[i] != '\n' {
372+
continue
373+
}
374+
if newlines++; newlines == limit.MaxLines {
375+
end = i
376+
break
377+
}
378+
}
379+
minEnd := start
380+
if start < len(text) {
381+
_, size := utf8.DecodeRuneInString(text[start:])
382+
minEnd = start + size
383+
}
384+
runeEnd := func(end int) int {
385+
end = max(end, minEnd)
386+
for end > minEnd && end < len(text) && !utf8.RuneStart(text[end]) {
387+
end--
388+
}
389+
return end
390+
}
391+
page := func(end int) map[string]any {
392+
out := map[string]any{"ok": true, "ref": ref, "tool_name": result.ToolName, "offset": start, "total_bytes": len(text), "content": text[start:end]}
393+
if result.IsError {
394+
out["is_error"] = true
395+
}
396+
if end < len(text) {
397+
out["next_offset"] = end
398+
}
399+
return out
400+
}
401+
fits := func(end int) bool { return len(contextlimit.MarshalString(page(end))) <= limit.MaxBytes }
402+
end = runeEnd(end)
403+
if !fits(end) {
404+
// Every shorter page carries next_offset, so its encoded size grows
405+
// with its end and the longest one that fits can be bisected.
406+
n := sort.Search(end-minEnd, func(n int) bool { return !fits(runeEnd(minEnd + n)) })
407+
end = runeEnd(minEnd + n - 1)
408+
}
409+
return page(end)
410+
}
411+
274412
func (p *HistoryProvider) ensureSessionVisible(ctx context.Context, sess SessionContext, sessionID string) error {
275413
_, allowed, err := visibleHistorySessions(ctx, p.sessions, sess)
276414
if err != nil {
@@ -644,7 +782,10 @@ type listSessionsArgs struct {
644782
type getMessagesArgs struct {
645783
Before string `json:"before,omitempty" jsonschema:"ISO 8601 timestamp cursor. When provided, returns messages created before this time."`
646784
Limit *int `json:"limit,omitempty" jsonschema:"Maximum number of messages to return. Default 30, max 100."`
785+
MaxBytes *int `json:"max_bytes,omitempty" jsonschema:"With ref: the most bytes of the result to return in this page. Default: as many as one tool result holds."`
647786
MessageID string `json:"message_id,omitempty" jsonschema:"Exact persisted message ID to resolve, such as a message_id returned in search_memory source_refs."`
787+
Offset *int `json:"offset,omitempty" jsonschema:"With ref: byte offset in the result to start from; pass the previous page's next_offset. Default 0."`
788+
Ref string `json:"ref,omitempty" jsonschema:"Tool result reference <message id>:<part index>, as cited by a summary (ref=...). Reads that earlier tool result's original text from the current session, one page at a time."`
648789
SessionID string `json:"session_id,omitempty" jsonschema:"Session ID to read. Defaults to the current session when omitted."`
649790
}
650791

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
package tools
2+
3+
import (
4+
"context"
5+
"os"
6+
"strings"
7+
"testing"
8+
9+
sdk "github.com/felinics/twilight/sdk"
10+
"github.com/google/uuid"
11+
12+
historyfrag "github.com/felinics/memoh/internal/agent/context/history"
13+
"github.com/felinics/memoh/internal/agent/toolexec"
14+
messagepkg "github.com/felinics/memoh/internal/chat/message"
15+
dbpkg "github.com/felinics/memoh/internal/db"
16+
"github.com/felinics/memoh/internal/db/dbtest"
17+
dbsqlc "github.com/felinics/memoh/internal/db/postgres/sqlc"
18+
postgresstore "github.com/felinics/memoh/internal/db/postgres/store"
19+
)
20+
21+
// A ref reads back, page by page, exactly the text the provider received for
22+
// a tool result persisted through the history write path, and a row of
23+
// another session of the same bot stays unreadable by ref.
24+
func TestPostgresHistoryToolResultRef(t *testing.T) {
25+
ctx := context.Background()
26+
dsn := strings.TrimSpace(os.Getenv("TEST_POSTGRES_DSN"))
27+
if dsn == "" {
28+
t.Skip("set TEST_POSTGRES_DSN to run the history ref PostgreSQL integration")
29+
}
30+
if os.Getenv("TEST_POSTGRES_BOOTSTRAP_SCHEMA") == "1" {
31+
if err := dbtest.MigratePostgresUp(dsn); err != nil {
32+
t.Fatalf("migrate: %v", err)
33+
}
34+
}
35+
pool, err := dbpkg.OpenPostgresDSN(ctx, dsn)
36+
if err != nil {
37+
t.Fatalf("open postgres: %v", err)
38+
}
39+
t.Cleanup(pool.Close)
40+
tx, err := pool.Begin(ctx)
41+
if err != nil {
42+
t.Fatalf("begin: %v", err)
43+
}
44+
t.Cleanup(func() { _ = tx.Rollback(ctx) })
45+
46+
userID, botID, currentID, otherID := uuid.NewString(), uuid.NewString(), uuid.NewString(), uuid.NewString()
47+
for _, fixture := range []struct {
48+
sql string
49+
args []any
50+
}{
51+
{`WITH u AS (INSERT INTO users (id, username, is_active) VALUES ($1::uuid, $1::text, true) RETURNING id)
52+
INSERT INTO team_members (user_id, role) SELECT id, 'admin' FROM u`, []any{userID}},
53+
{`INSERT INTO bots (id, owner_user_id, name) VALUES ($1::uuid, $2, $1::text)`, []any{botID, userID}},
54+
{`INSERT INTO bot_sessions (id, bot_id, channel_type) VALUES ($1, $3, 'local'), ($2, $3, 'local')`, []any{currentID, otherID, botID}},
55+
} {
56+
if _, err := tx.Exec(ctx, fixture.sql, fixture.args...); err != nil {
57+
t.Fatalf("fixture: %v", err)
58+
}
59+
}
60+
61+
messages := messagepkg.NewService(nil, postgresstore.NewQueries(dbsqlc.New(tx)))
62+
output, err := sdk.JSONOutput(map[string]any{"exit_code": 0, "stdout": strings.Repeat("line \"<quoted>\" & ünï 你好\n", 400)})
63+
if err != nil {
64+
t.Fatalf("encode output: %v", err)
65+
}
66+
content, err := historyfrag.MarshalStoredSDKMessage(sdk.Message{Role: sdk.MessageRoleTool, Content: []sdk.MessagePart{
67+
sdk.ToolResultPart{ToolCallID: "call-1", ToolName: "exec", Result: output},
68+
}})
69+
if err != nil {
70+
t.Fatalf("marshal tool row: %v", err)
71+
}
72+
rows := map[string]string{}
73+
for _, sessionID := range []string{currentID, otherID} {
74+
_, err := messages.Persist(ctx, messagepkg.PersistInput{BotID: botID, SessionID: sessionID, Role: "user", Content: []byte(`{"role":"user","content":"run it"}`)})
75+
row, toolErr := messages.Persist(ctx, messagepkg.PersistInput{BotID: botID, SessionID: sessionID, Role: "tool", Content: content})
76+
if err != nil || toolErr != nil {
77+
t.Fatalf("persist turn: %v, %v", err, toolErr)
78+
}
79+
rows[sessionID] = row.ID
80+
}
81+
82+
limit := ToolOutputLimit{MaxBytes: 4096, MaxLines: 64}
83+
tools, err := NewHistoryProvider(nil, nil, messages, nil, limit).Tools(ctx, SessionContext{BotID: botID, SessionID: currentID})
84+
if err != nil {
85+
t.Fatalf("Tools() error = %v", err)
86+
}
87+
getMessages := toolByNameForTest(t, WrapToolOutputLimits(tools, limit), ToolGetMessages())
88+
if pages := readToolResultRef(t, getMessages, rows[currentID]+":0", output.String()); pages < 2 {
89+
t.Fatalf("pages = %d, want the result split across pages", pages)
90+
}
91+
out, err := getMessages.Execute(&toolexec.ToolExecContext{Context: ctx, ToolName: getMessages.Name},
92+
toolexec.ArgumentsFromValue(map[string]any{"ref": rows[otherID] + ":0"}))
93+
if err == nil || !strings.Contains(err.Error(), "not in the current session") {
94+
t.Fatalf("ref into another session = %s, %v; want a current-session refusal", out.String(), err)
95+
}
96+
}

0 commit comments

Comments
 (0)
0