From 30d33977fcf4cbed7c73514c88c4c15598fea80b Mon Sep 17 00:00:00 2001 From: dwrz Date: Fri, 9 Oct 2026 15:15:21 +0000 Subject: [PATCH] Add job pkg metadata --- internal/job/job.go | 21 +++++++++++++++++---- internal/job/job_test.go | 39 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 4 deletions(-) diff --git a/internal/job/job.go b/internal/job/job.go index 8a7f236..24ce5da 100644 --- a/internal/job/job.go +++ b/internal/job/job.go @@ -8,6 +8,7 @@ package job import ( "context" "database/sql" + "encoding/json/jsontext" "encoding/json/v2" "errors" "fmt" @@ -254,7 +255,7 @@ func (j *Job) callTool( context.WithoutCancel(ctx), resultSaveTimeout, ) defer saveCancel() - if _, err := j.insertMessage(saveCtx, cid, message, update); err != nil { + if _, err := j.insertMessage(saveCtx, cid, message, nil, update); err != nil { log.ErrorContext( ctx, "failed to persist tool result", @@ -370,8 +371,8 @@ func (j *Job) createConversation( // Cron returns the parsed cron schedule. func (j *Job) Cron() cron.Schedule { return j.cron } -// insertMessage stores one message in a conversation, after every -// message already stored there. +// insertMessage stores a message and its optional inference metadata in the +// same row at the next sequence in a conversation. Nil metadata stores NULL. // If executionUpdate is non-nil, its status change and the message commit in // the same transaction. A nil update leaves execution statuses unchanged. // The caller is responsible for identifying the execution and choosing a @@ -380,12 +381,23 @@ func (j *Job) insertMessage( ctx context.Context, cid string, message llm.Message, + metadata *llm.Metadata, executionUpdate *db.UpdateExecutionStatusParams, ) (*db.Message, error) { content, err := json.Marshal(message) if err != nil { return nil, fmt.Errorf("marshal message: %w", err) } + var metadataJSON jsontext.Value + if metadata != nil { + encoded, err := json.Marshal( + metadata, json.FormatNilSliceAsNull(true), + ) + if err != nil { + return nil, fmt.Errorf("marshal metadata: %w", err) + } + metadataJSON = jsontext.Value(encoded) + } tx, err := j.db.BeginTx(ctx, nil) if err != nil { @@ -401,6 +413,7 @@ func (j *Job) insertMessage( m, err := q.InsertMessage(ctx, db.InsertMessageParams{ Content: string(content), ConversationId: cid, + Metadata: metadataJSON, Role: db.MessageRole(message.Role), Sequence: sequence, }) @@ -507,7 +520,7 @@ func (j *Job) Run(ctx context.Context) Result { // Persist the assistant message before acting on it, // so an interruption leaves a committed record. message := result.Message - tm, err = j.insertMessage(ctx, c.Id, message, nil) + tm, err = j.insertMessage(ctx, c.Id, message, &result.Metadata, nil) if err != nil { break } diff --git a/internal/job/job_test.go b/internal/job/job_test.go index 3d0ee0a..fa873a8 100644 --- a/internal/job/job_test.go +++ b/internal/job/job_test.go @@ -450,6 +450,32 @@ func TestRunPersistsToolLoop(t *testing.T) { if len(stored) != len(wantRoles) { t.Fatalf("stored messages = %d, want %d", len(stored), len(wantRoles)) } + for i, message := range stored { + if message.Role != db.MessageRoleAssistant { + if message.Metadata != nil { + t.Errorf("message %d metadata = %v, want nil", i, message.Metadata) + } + continue + } + if message.Metadata == nil { + t.Fatalf("assistant %d has no metadata", i) + } + var metadata llm.Metadata + if err := json.Unmarshal(message.Metadata, &metadata); err != nil { + t.Fatal(err) + } + if metadata.Provider == nil || *metadata.Provider != "test" || + metadata.RequestedModel == nil || *metadata.RequestedModel != "test-model" { + t.Errorf("assistant %d metadata = %s", i, string(message.Metadata)) + } + wantReason := "stop" + if i == 2 { + wantReason = "tool_calls" + } + if metadata.FinishReason == nil || *metadata.FinishReason != wantReason { + t.Errorf("assistant %d finish reason = %v", i, metadata.FinishReason) + } + } for i, want := range wantRoles { if stored[i].Sequence != i || stored[i].Role != want { t.Errorf( @@ -590,6 +616,7 @@ func TestCallTool(t *testing.T) { Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{call}, }, + &llm.Metadata{ID: new("assistant-completion")}, nil, ) if err != nil { @@ -644,7 +671,19 @@ func TestCallTool(t *testing.T) { if err != nil { t.Fatalf("json.Marshal: %v", err) } + wantMetadata := `{"choice_index":null,"created":null,` + + `"finish_reason":null,"id":"assistant-completion",` + + `"logprobs":null,"model":null,"provider":null,` + + `"requested_model":null,"service_tier":null,` + + `"system_fingerprint":null,"timings":null,"usage":null}` + if assistant.Metadata == nil || + string(assistant.Metadata) != wantMetadata { + t.Errorf("assistant metadata = %v", assistant.Metadata) + } last := stored[len(stored)-1] + if last.Metadata != nil { + t.Errorf("tool metadata = %v, want nil", last.Metadata) + } if last.Role != db.MessageRoleTool || last.Content != string(want) { t.Errorf(