diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index afa39fff..4ab88b53 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -168,6 +168,7 @@ func BuildRuntime(ctx context.Context, opts BootstrapOptions) (RuntimeBundle, er providerRegistry, contextBuilder, ) + runtimeSvc.SetSessionAssetStore(sessionStore) runtimeSvc.SetSkillsRegistry(buildSkillsRegistry(ctx, loader.BaseDir())) runtimeSvc.SetAutoCompactThresholdResolver(runtimeAutoCompactThresholdResolverFunc( func(ctx context.Context, cfg config.Config) (int, error) { diff --git a/internal/config/loader_test.go b/internal/config/loader_test.go index 2cee0104..a93c3d0f 100644 --- a/internal/config/loader_test.go +++ b/internal/config/loader_test.go @@ -1050,9 +1050,13 @@ func TestLoadCustomProvidersReadDirAndStatErrors(t *testing.T) { t.Fatalf("WriteFile() error = %v", err) } - if _, err := loadCustomProviders(baseDir); err == nil { + _, err := loadCustomProviders(baseDir) + if err == nil { t.Fatal("expected read providers dir error") } + if !strings.Contains(err.Error(), "read providers dir") { + t.Fatalf("expected read providers dir error, got %v", err) + } }) t.Run("provider yaml stat error", func(t *testing.T) { @@ -1061,17 +1065,42 @@ func TestLoadCustomProvidersReadDirAndStatErrors(t *testing.T) { if err := os.MkdirAll(providerDir, 0o755); err != nil { t.Fatalf("MkdirAll() error = %v", err) } - if err := os.Chmod(providerDir, 0o000); err != nil { - t.Fatalf("Chmod() error = %v", err) + // Windows 上 chmod(000) 不一定阻断访问，这里改为稳定的“provider.yaml 是目录”场景触发读取错误。 + if err := os.MkdirAll(filepath.Join(providerDir, customProviderConfigName), 0o755); err != nil { + t.Fatalf("MkdirAll(provider.yaml dir) error = %v", err) } - defer func() { _ = os.Chmod(providerDir, 0o755) }() - if _, err := loadCustomProviders(baseDir); err == nil { - t.Fatal("expected stat error") + _, err := loadCustomProviders(baseDir) + if err == nil { + t.Fatal("expected custom provider read error") + } + if !strings.Contains(err.Error(), "read") { + t.Fatalf("expected read error, got %v", err) } }) } +func TestLoadCustomProvidersReturnsEmptyWhenProvidersDirMissing(t *testing.T) { + t.Parallel() + + baseDir := t.TempDir() + providersPath := filepath.Join(baseDir, providersDirName) + if _, err := os.Stat(providersPath); !os.IsNotExist(err) { + t.Fatalf("expected providers dir to be missing, got %v", err) + } + + providers, err := loadCustomProviders(baseDir) + if err != nil { + t.Fatalf("loadCustomProviders() error = %v", err) + } + if len(providers) != 0 { + t.Fatalf("expected no custom providers, got %d", len(providers)) + } + if _, err := os.Stat(providersPath); !os.IsNotExist(err) { + t.Fatalf("expected providers dir to remain missing, got %v", err) + } +} + func TestLoadCustomProviderReadErrors(t *testing.T) { t.Run("missing provider yaml", func(t *testing.T) { providerDir := t.TempDir() diff --git a/internal/provider/openaicompat/chatcompletions/provider.go b/internal/provider/openaicompat/chatcompletions/provider.go index de9ea143..d92cce6d 100644 --- a/internal/provider/openaicompat/chatcompletions/provider.go +++ b/internal/provider/openaicompat/chatcompletions/provider.go @@ -38,7 +38,7 @@ func New(cfg provider.RuntimeConfig, client *http.Client) (*Provider, error) { // Generate 发起 SSE 流式生成请求。 // 流中途断连或协议错误时直接返回错误，由上层调用方决定重试策略。 func (p *Provider) Generate(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { - payload, err := BuildRequest(p.cfg, req) + payload, err := BuildRequest(ctx, p.cfg, req) if err != nil { return err } diff --git a/internal/provider/openaicompat/chatcompletions/provider_test.go b/internal/provider/openaicompat/chatcompletions/provider_test.go index 0b398dec..6a379d89 100644 --- a/internal/provider/openaicompat/chatcompletions/provider_test.go +++ b/internal/provider/openaicompat/chatcompletions/provider_test.go @@ -34,11 +34,21 @@ func TestNewAndBuildRequest(t *testing.T) { t.Run("build request variants", func(t *testing.T) { t.Parallel() - if _, err := BuildRequest(testCfg("https://api.example.com/v1", "", "test-key"), providertypes.GenerateRequest{}); err == nil || !strings.Contains(err.Error(), "model is empty") { + if _, err := BuildRequest(context.Background(), testCfg("https://api.example.com/v1", "", "test-key"), providertypes.GenerateRequest{}); err == nil || !strings.Contains(err.Error(), "model is empty") { t.Fatalf("expected model error, got %v", err) } + if _, err := BuildRequest(context.Background(), testCfg("https://api.example.com/v1", "gpt-4.1", "test-key"), providertypes.GenerateRequest{ + Messages: []providertypes.Message{ + { + Role: providertypes.RoleUser, + Parts: []providertypes.ContentPart{providertypes.NewSessionAssetImagePart("asset-without-reader", "image/png")}, + }, + }, + }); err == nil || !strings.Contains(err.Error(), "session_asset reader is not configured") { + t.Fatalf("expected BuildRequest conversion error, got %v", err) + } - payload, err := BuildRequest(testCfg("https://api.example.com/v1", "gpt-4.1", "test-key"), providertypes.GenerateRequest{ + payload, err := BuildRequest(context.Background(), testCfg("https://api.example.com/v1", "gpt-4.1", "test-key"), providertypes.GenerateRequest{ Model: "gpt-5.4", SystemPrompt: "system", Messages: []providertypes.Message{ @@ -61,7 +71,52 @@ func TestNewAndBuildRequest(t *testing.T) { t.Fatalf("unexpected tools: %+v", payload.Tools) } - fallback, err := BuildRequest(testCfg("https://api.example.com/v1", "gpt-4.1", "test-key"), providertypes.GenerateRequest{ + withSessionAsset, err := BuildRequest(context.Background(), testCfg("https://api.example.com/v1", "gpt-4.1", "test-key"), providertypes.GenerateRequest{ + Messages: []providertypes.Message{ + { + Role: providertypes.RoleUser, + Parts: []providertypes.ContentPart{ + providertypes.NewSessionAssetImagePart("asset-1", "image/png"), + }, + }, + }, + SessionAssetReader: stubSessionAssetReader{ + assets: map[string]stubSessionAsset{ + "asset-1": {data: []byte("img"), mime: "image/png"}, + }, + }, + }) + if err != nil { + t.Fatalf("BuildRequest() session asset error = %v", err) + } + converted, ok := withSessionAsset.Messages[0].Content.([]MessageContentPart) + if !ok || len(converted) != 1 || !strings.HasPrefix(converted[0].ImageURL.URL, "data:image/png;base64,") { + t.Fatalf("unexpected session asset conversion: %+v", withSessionAsset.Messages[0].Content) + } + + halfBudget := providertypes.MaxSessionAssetsTotalBytes / 2 + overBudgetSize := int(halfBudget + 1) + if _, err := BuildRequest(context.Background(), testCfg("https://api.example.com/v1", "gpt-4.1", "test-key"), providertypes.GenerateRequest{ + Messages: []providertypes.Message{ + { + Role: providertypes.RoleUser, + Parts: []providertypes.ContentPart{ + providertypes.NewSessionAssetImagePart("asset-1", "image/png"), + providertypes.NewSessionAssetImagePart("asset-2", "image/png"), + }, + }, + }, + SessionAssetReader: stubSessionAssetReader{ + assets: map[string]stubSessionAsset{ + "asset-1": {data: make([]byte, overBudgetSize), mime: "image/png"}, + "asset-2": {data: make([]byte, overBudgetSize), mime: "image/png"}, + }, + }, + }); err == nil || !strings.Contains(err.Error(), "session_asset total exceeds") { + t.Fatalf("expected session_asset total budget error, got %v", err) + } + + fallback, err := BuildRequest(context.Background(), testCfg("https://api.example.com/v1", "gpt-4.1", "test-key"), providertypes.GenerateRequest{ SystemPrompt: " ", Messages: []providertypes.Message{{Role: providertypes.RoleUser, Parts: []providertypes.ContentPart{providertypes.NewTextPart("hello")}}}, Tools: []providertypes.ToolSpec{}, @@ -77,7 +132,11 @@ func TestNewAndBuildRequest(t *testing.T) { t.Run("message conversion", func(t *testing.T) { t.Parallel() - user, err := ToOpenAIMessage(providertypes.Message{Role: providertypes.RoleUser, Parts: []providertypes.ContentPart{providertypes.NewTextPart("hello")}}) + user, err := ToOpenAIMessage( + context.Background(), + providertypes.Message{Role: providertypes.RoleUser, Parts: []providertypes.ContentPart{providertypes.NewTextPart("hello")}}, + nil, + ) if err != nil { t.Fatalf("ToOpenAIMessage() user error = %v", err) } @@ -85,13 +144,13 @@ func TestNewAndBuildRequest(t *testing.T) { t.Fatalf("unexpected user message: %+v", user) } - multiText, err := ToOpenAIMessage(providertypes.Message{ + multiText, err := ToOpenAIMessage(context.Background(), providertypes.Message{ Role: providertypes.RoleUser, Parts: []providertypes.ContentPart{ providertypes.NewTextPart("hello "), providertypes.NewTextPart("world"), }, - }) + }, nil) if err != nil { t.Fatalf("ToOpenAIMessage() multiText error = %v", err) } @@ -99,12 +158,12 @@ func TestNewAndBuildRequest(t *testing.T) { t.Fatalf("unexpected multiText content: %+v", multiText.Content) } - assistant, err := ToOpenAIMessage(providertypes.Message{ + assistant, err := ToOpenAIMessage(context.Background(), providertypes.Message{ Role: providertypes.RoleAssistant, ToolCalls: []providertypes.ToolCall{ {ID: "call_1", Name: "read_file", Arguments: `{"path":"main.go"}`}, }, - }) + }, nil) if err != nil { t.Fatalf("ToOpenAIMessage() assistant error = %v", err) } @@ -112,7 +171,11 @@ func TestNewAndBuildRequest(t *testing.T) { t.Fatalf("unexpected assistant message: %+v", assistant) } - tool, err := ToOpenAIMessage(providertypes.Message{Role: providertypes.RoleTool, ToolCallID: "call_1", Parts: []providertypes.ContentPart{providertypes.NewTextPart("result")}}) + tool, err := ToOpenAIMessage( + context.Background(), + providertypes.Message{Role: providertypes.RoleTool, ToolCallID: "call_1", Parts: []providertypes.ContentPart{providertypes.NewTextPart("result")}}, + nil, + ) if err != nil { t.Fatalf("ToOpenAIMessage() tool error = %v", err) } @@ -120,13 +183,13 @@ func TestNewAndBuildRequest(t *testing.T) { t.Fatalf("unexpected tool message: %+v", tool) } - multiModal, err := ToOpenAIMessage(providertypes.Message{ + multiModal, err := ToOpenAIMessage(context.Background(), providertypes.Message{ Role: providertypes.RoleUser, Parts: []providertypes.ContentPart{ providertypes.NewTextPart("look"), providertypes.NewRemoteImagePart("https://example.com/img.png"), }, - }) + }, nil) if err != nil { t.Fatalf("ToOpenAIMessage() multiModal error = %v", err) } @@ -135,38 +198,221 @@ func TestNewAndBuildRequest(t *testing.T) { t.Fatalf("unexpected multiModal message: %+v", multiModal) } - _, err = ToOpenAIMessage(providertypes.Message{ + _, err = ToOpenAIMessage(context.Background(), providertypes.Message{ + Role: providertypes.RoleUser, + Parts: []providertypes.ContentPart{ + providertypes.NewSessionAssetImagePart("asset-1", "image/png"), + }, + }, nil) + if err == nil || !strings.Contains(err.Error(), "session_asset reader is not configured") { + t.Fatalf("expected missing reader error, got %v", err) + } + _, err = ToOpenAIMessage(context.Background(), providertypes.Message{ + Role: providertypes.RoleUser, + Parts: []providertypes.ContentPart{ + providertypes.NewSessionAssetImagePart(" ", "image/png"), + }, + }, stubSessionAssetReader{}) + if err == nil || !strings.Contains(err.Error(), "missing asset id") { + t.Fatalf("expected missing asset id error, got %v", err) + } + + sessionAssetMsg, err := ToOpenAIMessage(context.Background(), providertypes.Message{ Role: providertypes.RoleUser, Parts: []providertypes.ContentPart{ providertypes.NewSessionAssetImagePart("asset-1", "image/png"), }, + }, stubSessionAssetReader{ + assets: map[string]stubSessionAsset{ + "asset-1": { + data: []byte("image-bytes"), + mime: "image/png", + }, + }, }) - if err == nil || !strings.Contains(err.Error(), "session_asset image is not supported") { - t.Fatalf("expected session_asset error, got %v", err) + if err != nil { + t.Fatalf("ToOpenAIMessage() session asset error = %v", err) + } + sessionParts, ok := sessionAssetMsg.Content.([]MessageContentPart) + if !ok || len(sessionParts) != 1 || sessionParts[0].Type != "image_url" || + !strings.HasPrefix(sessionParts[0].ImageURL.URL, "data:image/png;base64,") { + t.Fatalf("unexpected session asset message: %+v", sessionAssetMsg) } - _, err = ToOpenAIMessage(providertypes.Message{ + _, err = ToOpenAIMessage(context.Background(), providertypes.Message{ Role: providertypes.RoleUser, Parts: []providertypes.ContentPart{ - {Kind: providertypes.ContentPartImage, Image: &providertypes.ImagePart{SourceType: "unknown"}}, + providertypes.NewSessionAssetImagePart("asset-unsupported", "image/png"), + }, + }, stubSessionAssetReader{ + assets: map[string]stubSessionAsset{ + "asset-unsupported": { + data: []byte("data"), + mime: "application/octet-stream", + }, + }, + }) + if err == nil || !strings.Contains(err.Error(), "unsupported mime type") { + t.Fatalf("expected unsupported mime type error, got %v", err) + } + + _, err = ToOpenAIMessage(context.Background(), providertypes.Message{ + Role: providertypes.RoleUser, + Parts: []providertypes.ContentPart{ + providertypes.NewSessionAssetImagePart("asset-missing-mime", ""), + }, + }, stubSessionAssetReader{ + assets: map[string]stubSessionAsset{ + "asset-missing-mime": { + data: []byte("data"), + mime: " ", + }, + }, + }) + if err == nil || !strings.Contains(err.Error(), "missing mime type") { + t.Fatalf("expected missing mime type error, got %v", err) + } + + _, err = ToOpenAIMessage(context.Background(), providertypes.Message{ + Role: providertypes.RoleUser, + Parts: []providertypes.ContentPart{ + providertypes.NewSessionAssetImagePart("asset-big", "image/png"), + }, + }, stubSessionAssetReader{ + assets: map[string]stubSessionAsset{ + "asset-big": { + data: make([]byte, providertypes.MaxSessionAssetBytes+1), + mime: "image/png", + }, + }, + }) + if err == nil || !strings.Contains(err.Error(), "exceeds") { + t.Fatalf("expected oversized session asset error, got %v", err) + } + + canceledCtx, cancel := context.WithCancel(context.Background()) + cancel() + _, err = ToOpenAIMessage(canceledCtx, providertypes.Message{ + Role: providertypes.RoleUser, + Parts: []providertypes.ContentPart{ + providertypes.NewSessionAssetImagePart("asset-cancel", "image/png"), + }, + }, stubSessionAssetReader{ + openFunc: func(ctx context.Context, assetID string) (io.ReadCloser, string, error) { + if !errors.Is(ctx.Err(), context.Canceled) { + t.Fatalf("expected canceled ctx in reader, asset=%s err=%v", assetID, ctx.Err()) + } + return nil, "", ctx.Err() }, }) + if err == nil || !strings.Contains(err.Error(), "context canceled") { + t.Fatalf("expected context canceled error, got %v", err) + } + + _, err = ToOpenAIMessage(context.Background(), providertypes.Message{ + Role: providertypes.RoleUser, + Parts: []providertypes.ContentPart{ + {Kind: providertypes.ContentPartImage, Image: &providertypes.ImagePart{SourceType: "unknown"}}, + }, + }, nil) if err == nil || !strings.Contains(err.Error(), "unsupported source type") { t.Fatalf("expected unsupported image error, got %v", err) } - _, err = ToOpenAIMessage(providertypes.Message{ + _, err = ToOpenAIMessage(context.Background(), providertypes.Message{ Role: providertypes.RoleUser, Parts: []providertypes.ContentPart{ providertypes.NewRemoteImagePart(""), }, - }) + }, nil) if err == nil || !strings.Contains(err.Error(), "invalid message parts") { t.Fatalf("expected invalid parts error, got %v", err) } }) } +func TestResolveSessionAssetDataURLVariants(t *testing.T) { + t.Parallel() + + if _, _, err := resolveSessionAssetDataURL(context.Background(), stubSessionAssetReader{ + assets: map[string]stubSessionAsset{ + "asset-no-budget": {data: []byte("x"), mime: "image/png"}, + }, + }, &providertypes.AssetRef{ID: "asset-no-budget", MimeType: "image/png"}, 0); err == nil || + !strings.Contains(err.Error(), "session_asset total exceeds") { + t.Fatalf("expected no budget error, got %v", err) + } + + if _, _, err := resolveSessionAssetDataURL(nil, stubSessionAssetReader{ + assets: map[string]stubSessionAsset{ + "asset-nil-ctx": {data: []byte("x"), mime: "image/png"}, + }, + }, &providertypes.AssetRef{ID: "asset-nil-ctx", MimeType: "image/png"}, maxSessionAssetsTotalBytes); err != nil { + t.Fatalf("expected nil context to fallback to background, got %v", err) + } + + if _, _, err := resolveSessionAssetDataURL(context.Background(), stubSessionAssetReader{ + openFunc: func(ctx context.Context, assetID string) (io.ReadCloser, string, error) { + _ = ctx + _ = assetID + return nil, "", errors.New("open failed") + }, + }, &providertypes.AssetRef{ID: "asset-open-error", MimeType: "image/png"}, maxSessionAssetsTotalBytes); err == nil || !strings.Contains(err.Error(), "open session_asset") { + t.Fatalf("expected wrapped open error, got %v", err) + } + + if _, _, err := resolveSessionAssetDataURL(context.Background(), stubSessionAssetReader{ + openFunc: func(ctx context.Context, assetID string) (io.ReadCloser, string, error) { + _ = ctx + _ = assetID + return &failingReadCloser{err: errors.New("read failed")}, "image/png", nil + }, + }, &providertypes.AssetRef{ID: "asset-read-error", MimeType: "image/png"}, maxSessionAssetsTotalBytes); err == nil || !strings.Contains(err.Error(), "read session_asset") { + t.Fatalf("expected wrapped read error, got %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + if _, _, err := resolveSessionAssetDataURL(ctx, stubSessionAssetReader{ + openFunc: func(openCtx context.Context, assetID string) (io.ReadCloser, string, error) { + _ = openCtx + _ = assetID + return &cancelOnReadCloser{cancel: cancel}, "image/png", nil + }, + }, &providertypes.AssetRef{ID: "asset-cancel-after-read", MimeType: "image/png"}, maxSessionAssetsTotalBytes); err == nil || !strings.Contains(err.Error(), "context canceled") { + t.Fatalf("expected canceled context after read, got %v", err) + } + + if _, _, err := resolveSessionAssetDataURL(context.Background(), stubSessionAssetReader{ + assets: map[string]stubSessionAsset{ + "asset-empty": {data: []byte{}, mime: "image/png"}, + }, + }, &providertypes.AssetRef{ID: "asset-empty", MimeType: "image/png"}, maxSessionAssetsTotalBytes); err == nil || !strings.Contains(err.Error(), "is empty") { + t.Fatalf("expected empty asset error, got %v", err) + } + + if _, _, err := resolveSessionAssetDataURL(context.Background(), stubSessionAssetReader{ + assets: map[string]stubSessionAsset{ + "asset-over-total": {data: []byte("12345"), mime: "image/png"}, + }, + }, &providertypes.AssetRef{ID: "asset-over-total", MimeType: "image/png"}, 4); err == nil || + !strings.Contains(err.Error(), "session_asset total exceeds") { + t.Fatalf("expected total budget error, got %v", err) + } + + if _, _, err := toOpenAIMessageWithBudget(context.Background(), providertypes.Message{ + Role: providertypes.RoleUser, + Parts: []providertypes.ContentPart{ + providertypes.NewSessionAssetImagePart("asset-negative-budget", "image/png"), + }, + }, stubSessionAssetReader{ + assets: map[string]stubSessionAsset{ + "asset-negative-budget": {data: []byte("x"), mime: "image/png"}, + }, + }, -1); err == nil || !strings.Contains(err.Error(), "session_asset total exceeds") { + t.Fatalf("expected negative budget to fail as no budget, got %v", err) + } +} + func TestParseErrorVariants(t *testing.T) { t.Parallel() @@ -741,6 +987,23 @@ type cancelAfterDoneReader struct { read bool } +type cancelOnReadCloser struct { + cancel func() + read bool +} + +func (r *cancelOnReadCloser) Read(p []byte) (int, error) { + if r.read { + return 0, io.EOF + } + r.read = true + copy(p, []byte("img")) + r.cancel() + return len("img"), nil +} + +func (r *cancelOnReadCloser) Close() error { return nil } + func (r *cancelAfterDoneReader) Read(p []byte) (int, error) { if !r.read { r.read = true @@ -749,3 +1012,29 @@ func (r *cancelAfterDoneReader) Read(p []byte) (int, error) { r.cancel() return 0, r.err } + +type stubSessionAsset struct { + data []byte + mime string + err error +} + +type stubSessionAssetReader struct { + assets map[string]stubSessionAsset + openFunc func(ctx context.Context, assetID string) (io.ReadCloser, string, error) +} + +func (r stubSessionAssetReader) Open(ctx context.Context, assetID string) (io.ReadCloser, string, error) { + if r.openFunc != nil { + return r.openFunc(ctx, assetID) + } + _ = ctx + asset, ok := r.assets[assetID] + if !ok { + return nil, "", errors.New("asset not found") + } + if asset.err != nil { + return nil, "", asset.err + } + return io.NopCloser(strings.NewReader(string(asset.data))), asset.mime, nil +} diff --git a/internal/provider/openaicompat/chatcompletions/request.go b/internal/provider/openaicompat/chatcompletions/request.go index 283360f1..d37a658b 100644 --- a/internal/provider/openaicompat/chatcompletions/request.go +++ b/internal/provider/openaicompat/chatcompletions/request.go @@ -1,6 +1,8 @@ package chatcompletions import ( + "context" + "encoding/base64" "encoding/json" "errors" "fmt" @@ -13,9 +15,12 @@ import ( providertypes "neo-code/internal/provider/types" ) +const maxSessionAssetReadBytes = providertypes.MaxSessionAssetBytes +const maxSessionAssetsTotalBytes = providertypes.MaxSessionAssetsTotalBytes + // BuildRequest 将 provider.GenerateRequest 转换为 Chat Completions 请求结构。 // 模型优先取 req.Model，其次使用配置中的默认模型。 -func BuildRequest(cfg provider.RuntimeConfig, req providertypes.GenerateRequest) (Request, error) { +func BuildRequest(ctx context.Context, cfg provider.RuntimeConfig, req providertypes.GenerateRequest) (Request, error) { model := strings.TrimSpace(req.Model) if model == "" { model = strings.TrimSpace(cfg.DefaultModel) @@ -37,11 +42,19 @@ func BuildRequest(cfg provider.RuntimeConfig, req providertypes.GenerateRequest) }) } + var usedSessionAssetBytes int64 for _, message := range req.Messages { - msg, err := ToOpenAIMessage(message) + remainingSessionAssetBytes := maxSessionAssetsTotalBytes - usedSessionAssetBytes + msg, consumedBytes, err := toOpenAIMessageWithBudget( + ctx, + message, + req.SessionAssetReader, + remainingSessionAssetBytes, + ) if err != nil { return Request{}, err } + usedSessionAssetBytes += consumedBytes payload.Messages = append(payload.Messages, msg) } @@ -65,15 +78,30 @@ func BuildRequest(cfg provider.RuntimeConfig, req providertypes.GenerateRequest) } // ToOpenAIMessage 将通用 Message 转换为 OpenAI 协议消息格式。 -func ToOpenAIMessage(message providertypes.Message) (Message, error) { +func ToOpenAIMessage(ctx context.Context, message providertypes.Message, assetReader providertypes.SessionAssetReader) (Message, error) { + msg, _, err := toOpenAIMessageWithBudget(ctx, message, assetReader, maxSessionAssetsTotalBytes) + return msg, err +} + +// toOpenAIMessageWithBudget 将通用 Message 转换为 OpenAI 协议消息格式，并记录 session_asset 消耗字节数。 +func toOpenAIMessageWithBudget( + ctx context.Context, + message providertypes.Message, + assetReader providertypes.SessionAssetReader, + remainingAssetBudget int64, +) (Message, int64, error) { + if remainingAssetBudget < 0 { + remainingAssetBudget = 0 + } if err := providertypes.ValidateParts(message.Parts); err != nil { - return Message{}, fmt.Errorf("%sinvalid message parts: %w", shared.ErrorPrefix, err) + return Message{}, 0, fmt.Errorf("%sinvalid message parts: %w", shared.ErrorPrefix, err) } out := Message{ Role: message.Role, ToolCallID: message.ToolCallID, } + var usedAssetBytes int64 var hasImage bool for _, part := range message.Parts { @@ -112,9 +140,30 @@ func ToOpenAIMessage(message providertypes.Message) (Message, error) { }, }) case part.Image != nil && part.Image.SourceType == providertypes.ImageSourceSessionAsset: - return Message{}, errors.New("session_asset image is not supported in this phase") + if part.Image.Asset == nil || strings.TrimSpace(part.Image.Asset.ID) == "" { + return Message{}, 0, errors.New("session_asset image missing asset id") + } + if assetReader == nil { + return Message{}, 0, errors.New("session_asset reader is not configured") + } + imageURL, readBytes, err := resolveSessionAssetDataURL( + ctx, + assetReader, + part.Image.Asset, + remainingAssetBudget-usedAssetBytes, + ) + if err != nil { + return Message{}, 0, err + } + usedAssetBytes += readBytes + contentParts = append(contentParts, MessageContentPart{ + Type: "image_url", + ImageURL: &ImageURL{ + URL: imageURL, + }, + }) default: - return Message{}, errors.New("unsupported image part payload") + return Message{}, 0, errors.New("unsupported source type for image part") } } } @@ -137,7 +186,73 @@ func ToOpenAIMessage(message providertypes.Message) (Message, error) { } } - return out, nil + return out, usedAssetBytes, nil +} + +// resolveSessionAssetDataURL 读取会话附件并转换为可发送的 data URL，仅在请求阶段临时生成。 +func resolveSessionAssetDataURL( + ctx context.Context, + assetReader providertypes.SessionAssetReader, + asset *providertypes.AssetRef, + remainingBudget int64, +) (string, int64, error) { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return "", 0, err + } + if remainingBudget <= 0 { + return "", 0, fmt.Errorf( + "session_asset total exceeds %d bytes", + maxSessionAssetsTotalBytes, + ) + } + reader, mimeType, err := assetReader.Open(ctx, asset.ID) + if err != nil { + return "", 0, fmt.Errorf("open session_asset %q: %w", asset.ID, err) + } + defer func() { _ = reader.Close() }() + + readLimit := maxSessionAssetReadBytes + if remainingBudget < readLimit { + readLimit = remainingBudget + } + + data, err := io.ReadAll(io.LimitReader(reader, readLimit+1)) + if err != nil { + return "", 0, fmt.Errorf("read session_asset %q: %w", asset.ID, err) + } + if err := ctx.Err(); err != nil { + return "", 0, err + } + if int64(len(data)) > readLimit { + if readLimit < maxSessionAssetReadBytes { + return "", 0, fmt.Errorf( + "session_asset total exceeds %d bytes", + maxSessionAssetsTotalBytes, + ) + } + return "", 0, fmt.Errorf("session_asset %q exceeds %d bytes", asset.ID, maxSessionAssetReadBytes) + } + if len(data) == 0 { + return "", 0, fmt.Errorf("session_asset %q is empty", asset.ID) + } + + resolvedMime := strings.TrimSpace(mimeType) + if resolvedMime == "" { + resolvedMime = strings.TrimSpace(asset.MimeType) + } + normalizedMime := strings.ToLower(resolvedMime) + if normalizedMime == "" { + return "", 0, fmt.Errorf("session_asset %q missing mime type", asset.ID) + } + if !strings.HasPrefix(normalizedMime, "image/") { + return "", 0, fmt.Errorf("session_asset %q has unsupported mime type %q", asset.ID, resolvedMime) + } + + encoded := base64.StdEncoding.EncodeToString(data) + return fmt.Sprintf("data:%s;base64,%s", normalizedMime, encoded), int64(len(data)), nil } // ParseError 解析 HTTP 错误响应并包装为 ProviderError。 diff --git a/internal/provider/openaicompat/openaicompat_test.go b/internal/provider/openaicompat/openaicompat_test.go index 8190037e..7f002a83 100644 --- a/internal/provider/openaicompat/openaicompat_test.go +++ b/internal/provider/openaicompat/openaicompat_test.go @@ -159,7 +159,7 @@ func TestToOpenAIMessage_BasicMessage(t *testing.T) { Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hello world")}, } - result, err := chatcompletions.ToOpenAIMessage(msg) + result, err := chatcompletions.ToOpenAIMessage(context.Background(), msg, nil) if err != nil { t.Fatalf("ToOpenAIMessage() basic error = %v", err) } @@ -180,7 +180,7 @@ func TestToOpenAIMessage_ToolRoleMessage(t *testing.T) { Parts: []providertypes.ContentPart{providertypes.NewTextPart("result data")}, ToolCallID: "call_123", } - result, err := chatcompletions.ToOpenAIMessage(msg) + result, err := chatcompletions.ToOpenAIMessage(context.Background(), msg, nil) if err != nil { t.Fatalf("ToOpenAIMessage() tool role error = %v", err) } @@ -200,7 +200,7 @@ func TestToOpenAIMessage_AssistantWithToolCalls(t *testing.T) { {ID: "call_2", Name: "write_file", Arguments: `{"path":"test.go","content":"..."}`}, }, } - result, err := chatcompletions.ToOpenAIMessage(msg) + result, err := chatcompletions.ToOpenAIMessage(context.Background(), msg, nil) if err != nil { t.Fatalf("ToOpenAIMessage() assistant error = %v", err) } @@ -225,7 +225,7 @@ func TestToOpenAIMessage_EmptyToolCalls(t *testing.T) { t.Parallel() msg := providertypes.Message{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("test")}} - result, err := chatcompletions.ToOpenAIMessage(msg) + result, err := chatcompletions.ToOpenAIMessage(context.Background(), msg, nil) if err != nil { t.Fatalf("ToOpenAIMessage() empty tool calls error = %v", err) } @@ -298,7 +298,7 @@ func TestBuildRequest_EmptyModelReturnsError(t *testing.T) { client: &http.Client{}, } - _, buildErr := chatcompletions.BuildRequest(p.cfg, providertypes.GenerateRequest{}) + _, buildErr := chatcompletions.BuildRequest(context.Background(), p.cfg, providertypes.GenerateRequest{}) if buildErr == nil { t.Fatal("expected error for empty model") } @@ -315,7 +315,7 @@ func TestBuildRequest_FallsBackToConfigModel(t *testing.T) { t.Fatalf("New() error = %v", err) } - payload, err := chatcompletions.BuildRequest(p.cfg, providertypes.GenerateRequest{Messages: []providertypes.Message{{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hi")}}}}) + payload, err := chatcompletions.BuildRequest(context.Background(), p.cfg, providertypes.GenerateRequest{Messages: []providertypes.Message{{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hi")}}}}) if err != nil { t.Fatalf("buildRequest() error = %v", err) } @@ -332,7 +332,7 @@ func TestBuildRequest_RequestModelTakesPrecedence(t *testing.T) { t.Fatalf("New() error = %v", err) } - payload, err := chatcompletions.BuildRequest(p.cfg, providertypes.GenerateRequest{Model: "gpt-4-custom", Messages: []providertypes.Message{{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hi")}}}}) + payload, err := chatcompletions.BuildRequest(context.Background(), p.cfg, providertypes.GenerateRequest{Model: "gpt-4-custom", Messages: []providertypes.Message{{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hi")}}}}) if err != nil { t.Fatalf("buildRequest() error = %v", err) } @@ -349,7 +349,7 @@ func TestBuildRequest_NoSystemPrompt(t *testing.T) { t.Fatalf("New() error = %v", err) } - payload, err := chatcompletions.BuildRequest(p.cfg, providertypes.GenerateRequest{SystemPrompt: "", Messages: []providertypes.Message{{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hi")}}}}) + payload, err := chatcompletions.BuildRequest(context.Background(), p.cfg, providertypes.GenerateRequest{SystemPrompt: "", Messages: []providertypes.Message{{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hi")}}}}) if err != nil { t.Fatalf("buildRequest() error = %v", err) } @@ -368,7 +368,7 @@ func TestBuildRequest_NoTools(t *testing.T) { t.Fatalf("New() error = %v", err) } - payload, err := chatcompletions.BuildRequest(p.cfg, providertypes.GenerateRequest{Messages: []providertypes.Message{{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hi")}}}, Tools: nil}) + payload, err := chatcompletions.BuildRequest(context.Background(), p.cfg, providertypes.GenerateRequest{Messages: []providertypes.Message{{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hi")}}}, Tools: nil}) if err != nil { t.Fatalf("buildRequest() error = %v", err) } @@ -385,7 +385,7 @@ func TestBuildRequest_EmptyToolsSlice(t *testing.T) { t.Fatalf("New() error = %v", err) } - payload, err := chatcompletions.BuildRequest(p.cfg, providertypes.GenerateRequest{Messages: []providertypes.Message{{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hi")}}}, Tools: []providertypes.ToolSpec{}}) + payload, err := chatcompletions.BuildRequest(context.Background(), p.cfg, providertypes.GenerateRequest{Messages: []providertypes.Message{{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hi")}}}, Tools: []providertypes.ToolSpec{}}) if err != nil { t.Fatalf("buildRequest() error = %v", err) } @@ -402,7 +402,7 @@ func TestBuildRequest_MultipleTools(t *testing.T) { t.Fatalf("New() error = %v", err) } - payload, err := chatcompletions.BuildRequest(p.cfg, providertypes.GenerateRequest{ + payload, err := chatcompletions.BuildRequest(context.Background(), p.cfg, providertypes.GenerateRequest{ Messages: []providertypes.Message{{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("use tools")}}}, Tools: []providertypes.ToolSpec{ {Name: "tool_a", Description: "Tool A", Schema: map[string]any{"type": "object"}}, @@ -428,7 +428,7 @@ func TestBuildRequest_WhitespaceSystemPromptSkipped(t *testing.T) { t.Fatalf("New() error = %v", err) } - payload, err := chatcompletions.BuildRequest(p.cfg, providertypes.GenerateRequest{SystemPrompt: " ", Messages: []providertypes.Message{{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hi")}}}}) + payload, err := chatcompletions.BuildRequest(context.Background(), p.cfg, providertypes.GenerateRequest{SystemPrompt: " ", Messages: []providertypes.Message{{Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hi")}}}}) if err != nil { t.Fatalf("buildRequest() error = %v", err) } @@ -1191,7 +1191,7 @@ func TestBuildRequestIncludesSystemPromptToolsAndToolMessages(t *testing.T) { t.Fatalf("New() error = %v", err) } - payload, err := chatcompletions.BuildRequest(p.cfg, providertypes.GenerateRequest{ + payload, err := chatcompletions.BuildRequest(context.Background(), p.cfg, providertypes.GenerateRequest{ SystemPrompt: "system prompt", Messages: []providertypes.Message{ {Role: "user", Parts: []providertypes.ContentPart{providertypes.NewTextPart("hello")}}, diff --git a/internal/provider/types/asset_limits.go b/internal/provider/types/asset_limits.go new file mode 100644 index 00000000..5a315551 --- /dev/null +++ b/internal/provider/types/asset_limits.go @@ -0,0 +1,8 @@ +package types + +const ( + // MaxSessionAssetBytes 定义 session_asset 在读写链路中的统一大小上限（20 MiB）。 + MaxSessionAssetBytes int64 = 20 * 1024 * 1024 + // MaxSessionAssetsTotalBytes 定义单次请求允许携带的 session_asset 原始总字节上限（20 MiB）。 + MaxSessionAssetsTotalBytes int64 = 20 * 1024 * 1024 +) diff --git a/internal/provider/types/request.go b/internal/provider/types/request.go index 79ae4ba5..332d06a7 100644 --- a/internal/provider/types/request.go +++ b/internal/provider/types/request.go @@ -1,11 +1,22 @@ package types +import ( + "context" + "io" +) + +// SessionAssetReader 定义 provider 请求阶段读取会话附件内容的最小能力。 +type SessionAssetReader interface { + Open(ctx context.Context, assetID string) (io.ReadCloser, string, error) +} + // GenerateRequest 是 provider.Generate() 的请求参数。 type GenerateRequest struct { - Model string `json:"model"` - SystemPrompt string `json:"system_prompt"` - Messages []Message `json:"messages"` - Tools []ToolSpec `json:"tools,omitempty"` + Model string `json:"model"` + SystemPrompt string `json:"system_prompt"` + Messages []Message `json:"messages"` + Tools []ToolSpec `json:"tools,omitempty"` + SessionAssetReader SessionAssetReader `json:"-"` } // Usage 记录本次请求的 token 使用统计。 diff --git a/internal/runtime/run.go b/internal/runtime/run.go index 174eba9e..168580cf 100644 --- a/internal/runtime/run.go +++ b/internal/runtime/run.go @@ -272,10 +272,11 @@ func (s *Service) prepareTurnSnapshot(ctx context.Context, state *runState) (tur toolTimeout: time.Duration(cfg.ToolTimeoutSec) * time.Second, noProgressStreakLimit: limit, request: providertypes.GenerateRequest{ - Model: model, - SystemPrompt: systemPrompt, - Messages: builtContext.Messages, - Tools: toolSpecs, + Model: model, + SystemPrompt: systemPrompt, + Messages: builtContext.Messages, + Tools: toolSpecs, + SessionAssetReader: s.buildSessionAssetReader(ctx, state.session.ID), }, }, false, nil } diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index ae06687d..87c44d49 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -67,6 +67,7 @@ type AutoCompactThresholdResolver interface { type Service struct { configManager *config.Manager sessionStore agentsession.Store + sessionAssetStore agentsession.AssetStore toolManager tools.Manager providerFactory ProviderFactory contextBuilder agentcontext.Builder @@ -136,6 +137,11 @@ func (s *Service) SetMemoExtractor(extractor MemoExtractor) { s.memoExtractor = extractor } +// SetSessionAssetStore 设置会话附件存储实现，用于 provider 请求阶段读取 session_asset。 +func (s *Service) SetSessionAssetStore(store agentsession.AssetStore) { + s.sessionAssetStore = store +} + // SetSkillsRegistry 设置运行时可选的 skills registry，用于激活校验与上下文注入。 func (s *Service) SetSkillsRegistry(registry skills.Registry) { s.skillsRegistry = registry diff --git a/internal/runtime/session_assets.go b/internal/runtime/session_assets.go new file mode 100644 index 00000000..5e1c6bc4 --- /dev/null +++ b/internal/runtime/session_assets.go @@ -0,0 +1,49 @@ +package runtime + +import ( + "context" + "fmt" + "io" + "strings" + + providertypes "neo-code/internal/provider/types" + agentsession "neo-code/internal/session" +) + +// buildSessionAssetReader 按当前会话构造 provider 可用的附件读取器；缺失依赖时返回 nil。 +func (s *Service) buildSessionAssetReader(ctx context.Context, sessionID string) providertypes.SessionAssetReader { + if s.sessionAssetStore == nil { + return nil + } + id := strings.TrimSpace(sessionID) + if id == "" { + return nil + } + return runtimeSessionAssetReader{ + ctx: ctx, + store: s.sessionAssetStore, + sessionID: id, + } +} + +type runtimeSessionAssetReader struct { + ctx context.Context + store agentsession.AssetStore + sessionID string +} + +// Open 在会话范围内打开指定附件，供 provider 请求阶段读取并转换协议。 +func (r runtimeSessionAssetReader) Open(ctx context.Context, assetID string) (io.ReadCloser, string, error) { + openCtx := ctx + if openCtx == nil { + openCtx = r.ctx + } + if openCtx == nil { + openCtx = context.Background() + } + rc, meta, err := r.store.Open(openCtx, r.sessionID, strings.TrimSpace(assetID)) + if err != nil { + return nil, "", fmt.Errorf("runtime: open session asset %q: %w", assetID, err) + } + return rc, meta.MimeType, nil +} diff --git a/internal/runtime/session_assets_test.go b/internal/runtime/session_assets_test.go new file mode 100644 index 00000000..c6e718f9 --- /dev/null +++ b/internal/runtime/session_assets_test.go @@ -0,0 +1,168 @@ +package runtime + +import ( + "bytes" + "context" + "errors" + "io" + "strings" + "testing" + + agentsession "neo-code/internal/session" +) + +type stubAssetStore struct { + openFunc func(ctx context.Context, sessionID string, assetID string) (io.ReadCloser, agentsession.AssetMeta, error) +} + +func (s stubAssetStore) SaveAsset(ctx context.Context, sessionID string, r io.Reader, mimeType string) (agentsession.AssetMeta, error) { + _ = ctx + _ = sessionID + _ = r + _ = mimeType + return agentsession.AssetMeta{}, nil +} + +func (s stubAssetStore) Open(ctx context.Context, sessionID string, assetID string) (io.ReadCloser, agentsession.AssetMeta, error) { + return s.openFunc(ctx, sessionID, assetID) +} + +func (s stubAssetStore) Stat(ctx context.Context, sessionID string, assetID string) (agentsession.AssetMeta, error) { + _ = ctx + _ = sessionID + _ = assetID + return agentsession.AssetMeta{}, nil +} + +func TestBuildSessionAssetReaderReturnsNilWhenUnavailable(t *testing.T) { + t.Parallel() + + svc := &Service{} + if reader := svc.buildSessionAssetReader(context.Background(), "session-1"); reader != nil { + t.Fatalf("expected nil reader without asset store") + } + svc.SetSessionAssetStore(stubAssetStore{}) + if reader := svc.buildSessionAssetReader(context.Background(), " "); reader != nil { + t.Fatalf("expected nil reader for empty session id") + } +} + +func TestBuildSessionAssetReaderOpen(t *testing.T) { + t.Parallel() + + svc := &Service{} + svc.SetSessionAssetStore(stubAssetStore{ + openFunc: func(ctx context.Context, sessionID string, assetID string) (io.ReadCloser, agentsession.AssetMeta, error) { + if sessionID != "session-asset" || assetID != "asset-1" { + t.Fatalf("unexpected open args session=%q asset=%q", sessionID, assetID) + } + return io.NopCloser(bytes.NewReader([]byte("img"))), agentsession.AssetMeta{ + ID: "asset-1", + MimeType: "image/png", + Size: 3, + }, nil + }, + }) + + reader := svc.buildSessionAssetReader(context.Background(), "session-asset") + if reader == nil { + t.Fatalf("expected non-nil reader") + } + + rc, mime, err := reader.Open(context.Background(), "asset-1") + if err != nil { + t.Fatalf("reader.Open() error = %v", err) + } + defer func() { _ = rc.Close() }() + content, err := io.ReadAll(rc) + if err != nil { + t.Fatalf("ReadAll() error = %v", err) + } + if mime != "image/png" || string(content) != "img" { + t.Fatalf("unexpected open result mime=%q content=%q", mime, string(content)) + } +} + +func TestBuildSessionAssetReaderOpenErrorWrapped(t *testing.T) { + t.Parallel() + + svc := &Service{} + svc.SetSessionAssetStore(stubAssetStore{ + openFunc: func(ctx context.Context, sessionID string, assetID string) (io.ReadCloser, agentsession.AssetMeta, error) { + _ = ctx + _ = sessionID + _ = assetID + return nil, agentsession.AssetMeta{}, io.EOF + }, + }) + reader := svc.buildSessionAssetReader(context.Background(), "session-asset") + if reader == nil { + t.Fatalf("expected non-nil reader") + } + if _, _, err := reader.Open(context.Background(), "asset-1"); err == nil || !strings.Contains(err.Error(), "runtime: open session asset") { + t.Fatalf("expected wrapped open error, got %v", err) + } +} + +func TestBuildSessionAssetReaderOpenForwardsContext(t *testing.T) { + t.Parallel() + + svc := &Service{} + svc.SetSessionAssetStore(stubAssetStore{ + openFunc: func(ctx context.Context, sessionID string, assetID string) (io.ReadCloser, agentsession.AssetMeta, error) { + _ = sessionID + _ = assetID + if !errors.Is(ctx.Err(), context.Canceled) { + t.Fatalf("expected canceled context, got %v", ctx.Err()) + } + return nil, agentsession.AssetMeta{}, ctx.Err() + }, + }) + + runCtx, cancel := context.WithCancel(context.Background()) + cancel() + reader := svc.buildSessionAssetReader(runCtx, "session-asset") + if reader == nil { + t.Fatalf("expected non-nil reader") + } + if _, _, err := reader.Open(nil, "asset-ctx"); err == nil || !strings.Contains(err.Error(), "context canceled") { + t.Fatalf("expected context canceled error, got %v", err) + } +} + +func TestBuildSessionAssetReaderOpenFallsBackToBackgroundWhenNoContext(t *testing.T) { + t.Parallel() + + svc := &Service{} + svc.SetSessionAssetStore(stubAssetStore{ + openFunc: func(ctx context.Context, sessionID string, assetID string) (io.ReadCloser, agentsession.AssetMeta, error) { + if ctx == nil { + t.Fatalf("expected non-nil context") + } + if err := ctx.Err(); err != nil { + t.Fatalf("expected active context, got err=%v", err) + } + if sessionID != "session-bg" || assetID != "asset-bg" { + t.Fatalf("unexpected open args session=%q asset=%q", sessionID, assetID) + } + return io.NopCloser(bytes.NewReader([]byte("ok"))), agentsession.AssetMeta{ + ID: "asset-bg", + MimeType: "image/png", + Size: 2, + }, nil + }, + }) + + reader := svc.buildSessionAssetReader(nil, "session-bg") + if reader == nil { + t.Fatalf("expected non-nil reader") + } + rc, mime, err := reader.Open(nil, "asset-bg") + if err != nil { + t.Fatalf("reader.Open() error = %v", err) + } + defer func() { _ = rc.Close() }() + if mime != "image/png" { + t.Fatalf("expected mime image/png, got %q", mime) + } +} diff --git a/internal/session/asset_store.go b/internal/session/asset_store.go new file mode 100644 index 00000000..d63b4c16 --- /dev/null +++ b/internal/session/asset_store.go @@ -0,0 +1,87 @@ +package session + +import ( + "context" + "encoding/json" + "fmt" + "io" + "strings" +) + +// AssetMeta 描述会话附件最小元数据，用于 provider 请求阶段定位与发送。 +type AssetMeta struct { + ID string `json:"id"` + MimeType string `json:"mime_type"` + Size int64 `json:"size"` +} + +// AssetStore 定义会话级附件读写契约。 +type AssetStore interface { + SaveAsset(ctx context.Context, sessionID string, r io.Reader, mimeType string) (AssetMeta, error) + Open(ctx context.Context, sessionID string, assetID string) (io.ReadCloser, AssetMeta, error) + Stat(ctx context.Context, sessionID string, assetID string) (AssetMeta, error) +} + +type storedAssetMeta struct { + ID string `json:"id"` + MimeType string `json:"mime_type"` + Size int64 `json:"size"` +} + +// newAssetMeta 生成新的会话附件元数据，并校验 MIME 约束。 +func newAssetMeta(mimeType string) (AssetMeta, error) { + normalized := strings.ToLower(strings.TrimSpace(mimeType)) + if normalized == "" { + return AssetMeta{}, fmt.Errorf("session: asset mime type is empty") + } + if !strings.HasPrefix(normalized, "image/") { + return AssetMeta{}, fmt.Errorf("session: unsupported asset mime type %q", mimeType) + } + return AssetMeta{ + ID: NewID("asset"), + MimeType: normalized, + }, nil +} + +// decodeStoredAssetMeta 将磁盘上的附件元数据反序列化为运行时结构。 +func decodeStoredAssetMeta(data []byte) (AssetMeta, error) { + var stored storedAssetMeta + if err := json.Unmarshal(data, &stored); err != nil { + return AssetMeta{}, fmt.Errorf("session: decode asset meta: %w", err) + } + if err := validateStorageID("asset id", stored.ID); err != nil { + return AssetMeta{}, fmt.Errorf("session: %w", err) + } + normalizedMime := strings.ToLower(strings.TrimSpace(stored.MimeType)) + if normalizedMime == "" { + return AssetMeta{}, fmt.Errorf("session: asset meta mime_type is empty") + } + if !strings.HasPrefix(normalizedMime, "image/") { + return AssetMeta{}, fmt.Errorf("session: unsupported asset mime type %q", stored.MimeType) + } + return AssetMeta{ + ID: stored.ID, + MimeType: normalizedMime, + Size: stored.Size, + }, nil +} + +// encodeStoredAssetMeta 将运行时附件元数据编码为可持久化 JSON。 +func encodeStoredAssetMeta(meta AssetMeta) ([]byte, error) { + if err := validateStorageID("asset id", meta.ID); err != nil { + return nil, fmt.Errorf("session: %w", err) + } + normalizedMime := strings.ToLower(strings.TrimSpace(meta.MimeType)) + if normalizedMime == "" { + return nil, fmt.Errorf("session: asset mime type is empty") + } + payload, err := json.MarshalIndent(storedAssetMeta{ + ID: meta.ID, + MimeType: normalizedMime, + Size: meta.Size, + }, "", " ") + if err != nil { + return nil, fmt.Errorf("session: marshal asset meta: %w", err) + } + return append(payload, '\n'), nil +} diff --git a/internal/session/asset_store_test.go b/internal/session/asset_store_test.go new file mode 100644 index 00000000..7c2697a2 --- /dev/null +++ b/internal/session/asset_store_test.go @@ -0,0 +1,231 @@ +package session + +import ( + "bytes" + "context" + "errors" + "io" + "os" + "path/filepath" + "strings" + "testing" + + providertypes "neo-code/internal/provider/types" +) + +func TestJSONStoreSaveAssetOpenAndStat(t *testing.T) { + t.Parallel() + + store := NewJSONStore(t.TempDir(), t.TempDir()) + sessionID := "session_asset_round_trip" + payload := []byte("fake-image-bytes") + + meta, err := store.SaveAsset(context.Background(), sessionID, bytes.NewReader(payload), "image/png") + if err != nil { + t.Fatalf("SaveAsset() error = %v", err) + } + if meta.ID == "" || meta.MimeType != "image/png" || meta.Size != int64(len(payload)) { + t.Fatalf("unexpected saved meta: %+v", meta) + } + + statMeta, err := store.Stat(context.Background(), sessionID, meta.ID) + if err != nil { + t.Fatalf("Stat() error = %v", err) + } + if statMeta != meta { + t.Fatalf("expected stat meta %+v, got %+v", meta, statMeta) + } + + rc, openMeta, err := store.Open(context.Background(), sessionID, meta.ID) + if err != nil { + t.Fatalf("Open() error = %v", err) + } + defer func() { _ = rc.Close() }() + data, err := io.ReadAll(rc) + if err != nil { + t.Fatalf("ReadAll(open) error = %v", err) + } + if string(data) != string(payload) { + t.Fatalf("unexpected open payload: got %q want %q", string(data), string(payload)) + } + if openMeta != meta { + t.Fatalf("expected open meta %+v, got %+v", meta, openMeta) + } + + assetPath := store.assetPath(sessionID, meta.ID) + if _, err := os.Stat(assetPath); err != nil { + t.Fatalf("expected asset file %q exists, err=%v", assetPath, err) + } + if _, err := os.Stat(store.assetMetaPath(sessionID, meta.ID)); err != nil { + t.Fatalf("expected asset meta file exists, err=%v", err) + } +} + +func TestJSONStoreSaveAssetRejectsInvalidInput(t *testing.T) { + t.Parallel() + + store := NewJSONStore(t.TempDir(), t.TempDir()) + sessionID := "session_asset_invalid" + + if _, err := store.SaveAsset(context.Background(), sessionID, nil, "image/png"); err == nil { + t.Fatalf("expected nil reader error") + } + if _, err := store.SaveAsset(context.Background(), sessionID, strings.NewReader("x"), ""); err == nil { + t.Fatalf("expected empty mime error") + } + if _, err := store.SaveAsset(context.Background(), sessionID, strings.NewReader("x"), "text/plain"); err == nil { + t.Fatalf("expected unsupported mime error") + } + if _, err := store.SaveAsset(context.Background(), "../bad", strings.NewReader("x"), "image/png"); err == nil { + t.Fatalf("expected invalid session id error") + } +} + +func TestJSONStoreAssetOpenAndStatRejectInvalidID(t *testing.T) { + t.Parallel() + + store := NewJSONStore(t.TempDir(), t.TempDir()) + + if _, _, err := store.Open(context.Background(), "bad/session", "asset-ok"); err == nil { + t.Fatalf("expected invalid session id on open") + } + if _, _, err := store.Open(context.Background(), "session_ok", "../bad"); err == nil { + t.Fatalf("expected invalid asset id on open") + } + if _, err := store.Stat(context.Background(), "bad/session", "asset-ok"); err == nil { + t.Fatalf("expected invalid session id on stat") + } + if _, err := store.Stat(context.Background(), "session_ok", "../bad"); err == nil { + t.Fatalf("expected invalid asset id on stat") + } +} + +func TestJSONStoreAssetStoreRespectsCanceledContext(t *testing.T) { + t.Parallel() + + store := NewJSONStore(t.TempDir(), t.TempDir()) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + if _, err := store.SaveAsset(ctx, "session_ctx_cancel", strings.NewReader("x"), "image/png"); err == nil { + t.Fatalf("expected canceled SaveAsset error") + } + if _, _, err := store.Open(ctx, "session_ctx_cancel", "asset_x"); err == nil { + t.Fatalf("expected canceled Open error") + } + if _, err := store.Stat(ctx, "session_ctx_cancel", "asset_x"); err == nil { + t.Fatalf("expected canceled Stat error") + } +} + +func TestJSONStoreSaveAssetRejectsOversizedPayload(t *testing.T) { + t.Parallel() + + store := NewJSONStore(t.TempDir(), t.TempDir()) + oversized := bytes.NewReader(make([]byte, providertypes.MaxSessionAssetBytes+1)) + + if _, err := store.SaveAsset(context.Background(), "session_oversize", oversized, "image/png"); err == nil || + !strings.Contains(err.Error(), "asset size exceeds") { + t.Fatalf("expected oversized payload rejection, got %v", err) + } +} + +func TestDecodeStoredAssetMetaRejectsNonImageMIME(t *testing.T) { + t.Parallel() + + _, err := decodeStoredAssetMeta([]byte(`{"id":"asset_ok","mime_type":"text/plain","size":1}`)) + if err == nil || !strings.Contains(err.Error(), "unsupported asset mime type") { + t.Fatalf("expected non-image mime rejection, got %v", err) + } +} + +func TestDecodeAndEncodeStoredAssetMetaValidation(t *testing.T) { + t.Parallel() + + if _, err := decodeStoredAssetMeta([]byte(`{`)); err == nil || !strings.Contains(err.Error(), "decode asset meta") { + t.Fatalf("expected decode error, got %v", err) + } + if _, err := decodeStoredAssetMeta([]byte(`{"id":"bad/asset","mime_type":"image/png","size":1}`)); err == nil || + !strings.Contains(err.Error(), "unsupported characters") { + t.Fatalf("expected invalid asset id error, got %v", err) + } + if _, err := decodeStoredAssetMeta([]byte(`{"id":"asset_ok","mime_type":" ","size":1}`)); err == nil || + !strings.Contains(err.Error(), "mime_type is empty") { + t.Fatalf("expected empty mime_type error, got %v", err) + } + + if _, err := encodeStoredAssetMeta(AssetMeta{ID: "bad/asset", MimeType: "image/png", Size: 1}); err == nil || + !strings.Contains(err.Error(), "unsupported characters") { + t.Fatalf("expected invalid asset id encode error, got %v", err) + } + if _, err := encodeStoredAssetMeta(AssetMeta{ID: "asset_ok", MimeType: " ", Size: 1}); err == nil || + !strings.Contains(err.Error(), "asset mime type is empty") { + t.Fatalf("expected empty mime encode error, got %v", err) + } +} + +func TestJSONStoreSaveAssetFailurePaths(t *testing.T) { + t.Parallel() + + t.Run("create assets dir failed", func(t *testing.T) { + t.Parallel() + + baseDir := t.TempDir() + workspaceRoot := t.TempDir() + store := NewJSONStore(baseDir, workspaceRoot) + + assetsDir := store.assetsDir("session_assets_dir_fail") + if err := os.MkdirAll(filepath.Dir(assetsDir), 0o755); err != nil { + t.Fatalf("mkdir parent: %v", err) + } + if err := os.WriteFile(assetsDir, []byte("blocked"), 0o644); err != nil { + t.Fatalf("write blocker file: %v", err) + } + + if _, err := store.SaveAsset(context.Background(), "session_assets_dir_fail", strings.NewReader("x"), "image/png"); err == nil || + !strings.Contains(err.Error(), "create assets dir") { + t.Fatalf("expected create assets dir error, got %v", err) + } + }) + + t.Run("copy temp asset failed", func(t *testing.T) { + t.Parallel() + + store := NewJSONStore(t.TempDir(), t.TempDir()) + if _, err := store.SaveAsset(context.Background(), "session_copy_fail", failingReader{}, "image/png"); err == nil || + !strings.Contains(err.Error(), "write temp asset") { + t.Fatalf("expected write temp asset error, got %v", err) + } + }) +} + +func TestJSONStoreOpenAndStatMissingStoredFiles(t *testing.T) { + t.Parallel() + + store := NewJSONStore(t.TempDir(), t.TempDir()) + sessionID := "session_missing_files" + meta, err := store.SaveAsset(context.Background(), sessionID, strings.NewReader("img"), "image/png") + if err != nil { + t.Fatalf("save seed asset: %v", err) + } + + if err := os.Remove(store.assetPath(sessionID, meta.ID)); err != nil { + t.Fatalf("remove asset file: %v", err) + } + if _, _, err := store.Open(context.Background(), sessionID, meta.ID); err == nil { + t.Fatalf("expected open failure when asset binary is missing") + } + + if err := os.Remove(store.assetMetaPath(sessionID, meta.ID)); err != nil { + t.Fatalf("remove asset meta file: %v", err) + } + if _, err := store.Stat(context.Background(), sessionID, meta.ID); err == nil { + t.Fatalf("expected stat failure when asset meta is missing") + } +} + +type failingReader struct{} + +func (failingReader) Read(_ []byte) (int, error) { + return 0, errors.New("read failure") +} diff --git a/internal/session/skill_activation_test.go b/internal/session/skill_activation_test.go index 10be3db3..19704a6b 100644 --- a/internal/session/skill_activation_test.go +++ b/internal/session/skill_activation_test.go @@ -3,7 +3,6 @@ package session import ( "context" "os" - "path/filepath" "strings" "testing" "time" @@ -75,7 +74,7 @@ func TestJSONStoreSaveLoadRoundTripActivatedSkills(t *testing.T) { t.Fatalf("expected normalized loaded activations, got %+v", got) } - rawPath := filepath.Join(sessionDirectory(baseDir, workspaceRoot), session.ID+".json") + rawPath := sessionFilePathForTest(baseDir, workspaceRoot, session.ID) raw, err := os.ReadFile(rawPath) if err != nil { t.Fatalf("read saved session: %v", err) @@ -92,7 +91,7 @@ func TestJSONStoreLoadAllowsMissingActivatedSkillsField(t *testing.T) { workspaceRoot := t.TempDir() store := NewJSONStore(baseDir, workspaceRoot) - mustWriteSessionFile(t, filepath.Join(sessionDirectory(baseDir, workspaceRoot), "no-activated-skills.json"), strings.Join([]string{ + mustWriteSessionFile(t, sessionFilePathForTest(baseDir, workspaceRoot, "no-activated-skills"), strings.Join([]string{ `{`, ` "schema_version": 2,`, ` "id": "no-activated-skills",`, diff --git a/internal/session/store.go b/internal/session/store.go index 8824af65..147a3527 100644 --- a/internal/session/store.go +++ b/internal/session/store.go @@ -5,8 +5,10 @@ import ( "encoding/json" "errors" "fmt" + "io" "os" "path/filepath" + "regexp" "sort" "strings" "sync" @@ -17,6 +19,13 @@ import ( const sessionsDirName = "sessions" +const ( + sessionFileName = "session.json" + assetsDirName = "assets" +) + +var storageIDPattern = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_-]{0,127}$`) + // Session 表示单个会话的持久化模型，包含基础元数据与消息历史。 // Provider / Model 用于在 compact 等流程中优先复用会话最近一次成功运行的模型配置。 type Session struct { @@ -83,6 +92,9 @@ func (s *JSONStore) Save(ctx context.Context, session *Session) error { if err := validateSessionSchema(*session); err != nil { return err } + if err := validateStorageID("session id", session.ID); err != nil { + return fmt.Errorf("session: %w", err) + } session.TaskState = normalizeAndClampTaskState(session.TaskState) session.ActivatedSkills = normalizeSkillActivations(session.ActivatedSkills) @@ -109,15 +121,14 @@ func (s *JSONStore) Save(ctx context.Context, session *Session) error { payload = append(payload, '\n') target := s.filePath(session.ID) - temp := target + ".tmp" - if err := os.WriteFile(temp, payload, 0o644); err != nil { - return fmt.Errorf("session: write temp session: %w", err) + if err := ensurePathWithinBase(s.baseDir, target); err != nil { + return fmt.Errorf("session: resolve session file path: %w", err) } - if err := os.Remove(target); err != nil && !errors.Is(err, os.ErrNotExist) { - return fmt.Errorf("session: replace session file: %w", err) + if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil { + return fmt.Errorf("session: create session dir: %w", err) } - if err := os.Rename(temp, target); err != nil { - return fmt.Errorf("session: commit session file: %w", err) + if err := writeFileAtomically(target, "session-*.tmp", payload, 0o644); err != nil { + return err } return nil @@ -129,10 +140,19 @@ func (s *JSONStore) Load(ctx context.Context, id string) (Session, error) { return Session{}, err } + if err := validateStorageID("session id", id); err != nil { + return Session{}, fmt.Errorf("session: %w", err) + } + s.mu.RLock() defer s.mu.RUnlock() - data, err := os.ReadFile(s.filePath(id)) + target := s.filePath(id) + if err := ensurePathWithinBase(s.baseDir, target); err != nil { + return Session{}, fmt.Errorf("session: resolve session file path: %w", err) + } + + data, err := os.ReadFile(target) if err != nil { return Session{}, err } @@ -164,7 +184,7 @@ func (s *JSONStore) ListSummaries(ctx context.Context) ([]Summary, error) { summaries := make([]Summary, 0, len(entries)) for _, entry := range entries { - if entry.IsDir() || filepath.Ext(entry.Name()) != ".json" { + if !entry.IsDir() { continue } @@ -174,7 +194,12 @@ func (s *JSONStore) ListSummaries(ctx context.Context) ([]Summary, error) { default: } - data, readErr := os.ReadFile(filepath.Join(s.baseDir, entry.Name())) + target := filepath.Join(s.baseDir, entry.Name(), sessionFileName) + if err := ensurePathWithinBase(s.baseDir, target); err != nil { + continue + } + + data, readErr := os.ReadFile(target) if readErr != nil { continue } @@ -198,7 +223,170 @@ func (s *JSONStore) ListSummaries(ctx context.Context) ([]Summary, error) { // filePath 生成会话 ID 对应的 JSON 文件路径。 func (s *JSONStore) filePath(id string) string { - return filepath.Join(s.baseDir, id+".json") + return filepath.Join(s.sessionDir(id), sessionFileName) +} + +// sessionDir 返回指定会话在当前工作区分桶下的目录路径。 +func (s *JSONStore) sessionDir(id string) string { + return filepath.Join(s.baseDir, id) +} + +// assetsDir 返回指定会话附件目录路径。 +func (s *JSONStore) assetsDir(sessionID string) string { + return filepath.Join(s.sessionDir(sessionID), assetsDirName) +} + +// assetPath 返回指定会话附件二进制文件路径。 +func (s *JSONStore) assetPath(sessionID string, assetID string) string { + return filepath.Join(s.assetsDir(sessionID), assetID+".bin") +} + +// assetMetaPath 返回指定会话附件元数据文件路径。 +func (s *JSONStore) assetMetaPath(sessionID string, assetID string) string { + return filepath.Join(s.assetsDir(sessionID), assetID+".json") +} + +// SaveAsset 将会话附件二进制内容写入当前工作区会话目录，并返回附件元数据。 +func (s *JSONStore) SaveAsset(ctx context.Context, sessionID string, r io.Reader, mimeType string) (AssetMeta, error) { + if err := ctx.Err(); err != nil { + return AssetMeta{}, err + } + if r == nil { + return AssetMeta{}, errors.New("session: asset reader is nil") + } + if err := validateStorageID("session id", sessionID); err != nil { + return AssetMeta{}, fmt.Errorf("session: %w", err) + } + + meta, err := newAssetMeta(mimeType) + if err != nil { + return AssetMeta{}, err + } + + s.mu.Lock() + defer s.mu.Unlock() + + assetDir := s.assetsDir(sessionID) + if err := ensurePathWithinBase(s.baseDir, assetDir); err != nil { + return AssetMeta{}, fmt.Errorf("session: resolve assets dir path: %w", err) + } + if err := os.MkdirAll(assetDir, 0o755); err != nil { + return AssetMeta{}, fmt.Errorf("session: create assets dir: %w", err) + } + + target := s.assetPath(sessionID, meta.ID) + if err := ensurePathWithinBase(s.baseDir, target); err != nil { + return AssetMeta{}, fmt.Errorf("session: resolve asset file path: %w", err) + } + tempFile, tempPath, err := createTempFile(assetDir, "asset-*.tmp", "create temp asset") + if err != nil { + return AssetMeta{}, err + } + + written, copyErr := io.Copy(tempFile, io.LimitReader(r, providertypes.MaxSessionAssetBytes+1)) + syncErr := tempFile.Sync() + closeErr := tempFile.Close() + if copyErr != nil { + _ = os.Remove(tempPath) + return AssetMeta{}, fmt.Errorf("session: write temp asset: %w", copyErr) + } + if written > providertypes.MaxSessionAssetBytes { + _ = os.Remove(tempPath) + return AssetMeta{}, fmt.Errorf("session: asset size exceeds %d bytes", providertypes.MaxSessionAssetBytes) + } + if syncErr != nil { + _ = os.Remove(tempPath) + return AssetMeta{}, fmt.Errorf("session: sync temp asset: %w", syncErr) + } + if closeErr != nil { + _ = os.Remove(tempPath) + return AssetMeta{}, fmt.Errorf("session: close temp asset: %w", closeErr) + } + + meta.Size = written + if err := replaceFileWithTemp(tempPath, target, "asset file"); err != nil { + _ = os.Remove(tempPath) + return AssetMeta{}, err + } + + metaData, err := encodeStoredAssetMeta(meta) + if err != nil { + _ = os.Remove(target) + return AssetMeta{}, err + } + metaTarget := s.assetMetaPath(sessionID, meta.ID) + if err := ensurePathWithinBase(s.baseDir, metaTarget); err != nil { + _ = os.Remove(target) + return AssetMeta{}, fmt.Errorf("session: resolve asset meta file path: %w", err) + } + if err := writeFileAtomically(metaTarget, "asset-meta-*.tmp", metaData, 0o644); err != nil { + _ = os.Remove(target) + return AssetMeta{}, err + } + + return meta, nil +} + +// Open 读取会话附件二进制内容并返回可关闭流与附件元数据。 +func (s *JSONStore) Open(ctx context.Context, sessionID string, assetID string) (io.ReadCloser, AssetMeta, error) { + if err := ctx.Err(); err != nil { + return nil, AssetMeta{}, err + } + if err := validateStorageID("session id", sessionID); err != nil { + return nil, AssetMeta{}, fmt.Errorf("session: %w", err) + } + if err := validateStorageID("asset id", assetID); err != nil { + return nil, AssetMeta{}, fmt.Errorf("session: %w", err) + } + + s.mu.RLock() + defer s.mu.RUnlock() + + meta, err := s.statUnlocked(sessionID, assetID) + if err != nil { + return nil, AssetMeta{}, err + } + + target := s.assetPath(sessionID, assetID) + if err := ensurePathWithinBase(s.baseDir, target); err != nil { + return nil, AssetMeta{}, fmt.Errorf("session: resolve asset file path: %w", err) + } + file, err := os.Open(target) + if err != nil { + return nil, AssetMeta{}, err + } + return file, meta, nil +} + +// Stat 返回会话附件的元数据而不读取实际内容。 +func (s *JSONStore) Stat(ctx context.Context, sessionID string, assetID string) (AssetMeta, error) { + if err := ctx.Err(); err != nil { + return AssetMeta{}, err + } + if err := validateStorageID("session id", sessionID); err != nil { + return AssetMeta{}, fmt.Errorf("session: %w", err) + } + if err := validateStorageID("asset id", assetID); err != nil { + return AssetMeta{}, fmt.Errorf("session: %w", err) + } + + s.mu.RLock() + defer s.mu.RUnlock() + + return s.statUnlocked(sessionID, assetID) +} + +// statUnlocked 在调用方已持有读锁时读取附件元数据，避免重复加锁导致死锁风险。 +func (s *JSONStore) statUnlocked(sessionID string, assetID string) (AssetMeta, error) { + target := s.assetMetaPath(sessionID, assetID) + if err := ensurePathWithinBase(s.baseDir, target); err != nil { + return AssetMeta{}, fmt.Errorf("session: resolve asset meta file path: %w", err) + } + data, err := os.ReadFile(target) + if err != nil { + return AssetMeta{}, err + } + return decodeStoredAssetMeta(data) } // New 创建一个默认标题策略的新会话对象。 @@ -350,3 +538,97 @@ func decodeStoredSummary(data []byte) (Summary, error) { UpdatedAt: stored.UpdatedAt, }, nil } + +// validateStorageID 校验会话/附件 ID，避免路径穿越和非法文件名。 +func validateStorageID(label string, id string) error { + trimmed := strings.TrimSpace(id) + if trimmed == "" { + return fmt.Errorf("%s is empty", label) + } + if !storageIDPattern.MatchString(trimmed) { + return fmt.Errorf("%s %q contains unsupported characters", label, id) + } + return nil +} + +// ensurePathWithinBase 校验目标路径在给定基目录内，作为 ID 白名单之外的二次路径约束。 +func ensurePathWithinBase(baseDir string, target string) error { + baseAbs, err := filepath.Abs(baseDir) + if err != nil { + return fmt.Errorf("resolve base dir %q: %w", baseDir, err) + } + targetAbs, err := filepath.Abs(target) + if err != nil { + return fmt.Errorf("resolve target path %q: %w", target, err) + } + rel, err := filepath.Rel(baseAbs, targetAbs) + if err != nil { + return fmt.Errorf("compute relative path %q -> %q: %w", baseAbs, targetAbs, err) + } + if rel == "." { + return nil + } + if !filepath.IsLocal(rel) { + return fmt.Errorf("target path %q escapes base dir %q", targetAbs, baseAbs) + } + return nil +} + +// createTempFile 在目标目录创建唯一临时文件，避免固定 *.tmp 命名在并发场景下冲突。 +func createTempFile(dir string, pattern string, op string) (*os.File, string, error) { + file, err := os.CreateTemp(dir, pattern) + if err != nil { + return nil, "", fmt.Errorf("session: %s: %w", op, err) + } + if err := ensurePathWithinBase(dir, file.Name()); err != nil { + _ = file.Close() + _ = os.Remove(file.Name()) + return nil, "", fmt.Errorf("session: %s: %w", op, err) + } + return file, file.Name(), nil +} + +// replaceFileWithTemp 使用原子重命名替换目标文件，兼容 Windows 需要先删除旧文件的行为。 +func replaceFileWithTemp(tempPath string, target string, label string) error { + if err := os.Remove(target); err != nil && !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("session: replace %s: %w", label, err) + } + if err := os.Rename(tempPath, target); err != nil { + return fmt.Errorf("session: commit %s: %w", label, err) + } + return nil +} + +// writeFileAtomically 将字节数据写入唯一临时文件并原子替换目标文件，避免中间态文件暴露。 +func writeFileAtomically(target string, tempPattern string, payload []byte, perm os.FileMode) error { + dir := filepath.Dir(target) + tempFile, tempPath, err := createTempFile(dir, tempPattern, "create temp file") + if err != nil { + return err + } + + _, writeErr := tempFile.Write(payload) + syncErr := tempFile.Sync() + closeErr := tempFile.Close() + if writeErr != nil { + _ = os.Remove(tempPath) + return fmt.Errorf("session: write temp file: %w", writeErr) + } + if syncErr != nil { + _ = os.Remove(tempPath) + return fmt.Errorf("session: sync temp file: %w", syncErr) + } + if closeErr != nil { + _ = os.Remove(tempPath) + return fmt.Errorf("session: close temp file: %w", closeErr) + } + if err := os.Chmod(tempPath, perm); err != nil { + _ = os.Remove(tempPath) + return fmt.Errorf("session: chmod temp file: %w", err) + } + if err := replaceFileWithTemp(tempPath, target, "file"); err != nil { + _ = os.Remove(tempPath) + return err + } + return nil +} diff --git a/internal/session/store_test.go b/internal/session/store_test.go index 02c2b651..54d92baf 100644 --- a/internal/session/store_test.go +++ b/internal/session/store_test.go @@ -68,7 +68,7 @@ func TestJSONStoreSaveLoadAndListSummaries(t *testing.T) { t.Fatalf("unexpected loaded messages: %+v", loaded.Messages) } - rawPath := filepath.Join(sessionDirectory(baseDir, workspaceRoot), newer.ID+".json") + rawPath := sessionFilePathForTest(baseDir, workspaceRoot, newer.ID) raw, err := os.ReadFile(rawPath) if err != nil { t.Fatalf("read saved session: %v", err) @@ -77,7 +77,7 @@ func TestJSONStoreSaveLoadAndListSummaries(t *testing.T) { t.Fatalf("expected persisted session file to include workdir, got:\n%s", string(raw)) } - mustWriteSessionFile(t, filepath.Join(sessionDirectory(baseDir, workspaceRoot), "invalid.json"), "{invalid") + mustWriteSessionFile(t, sessionFilePathForTest(baseDir, workspaceRoot, "invalid"), "{invalid") if err := os.MkdirAll(filepath.Join(sessionDirectory(baseDir, workspaceRoot), "directory"), 0o755); err != nil { t.Fatalf("mkdir stray directory: %v", err) } @@ -215,6 +215,9 @@ func TestJSONStoreErrors(t *testing.T) { if _, err := store.ListSummaries(cancelledCtx); err == nil { t.Fatalf("expected cancelled list to fail") } + if _, err := store.Load(context.Background(), " "); err == nil || !strings.Contains(err.Error(), "session id is empty") { + t.Fatalf("expected empty session id load error, got %v", err) + } } func TestJSONStoreCorruptedSessionBehaviors(t *testing.T) { @@ -236,7 +239,7 @@ func TestJSONStoreCorruptedSessionBehaviors(t *testing.T) { t.Fatalf("Save valid session: %v", err) } - mustWriteSessionFile(t, filepath.Join(sessionDirectory(baseDir, workspaceRoot), "broken.json"), "{broken") + mustWriteSessionFile(t, sessionFilePathForTest(baseDir, workspaceRoot, "broken"), "{broken") _, err := store.Load(context.Background(), "broken") if err == nil || !strings.Contains(err.Error(), "decode session broken") { @@ -280,7 +283,7 @@ func TestJSONStoreSaveReplaceFailureWhenTargetIsNonEmptyDirectory(t *testing.T) baseDir := t.TempDir() workspaceRoot := t.TempDir() store := NewJSONStore(baseDir, workspaceRoot) - targetDir := filepath.Join(sessionDirectory(baseDir, workspaceRoot), "blocked.json") + targetDir := sessionFilePathForTest(baseDir, workspaceRoot, "blocked") if err := os.MkdirAll(targetDir, 0o755); err != nil { t.Fatalf("mkdir target dir: %v", err) } @@ -295,7 +298,7 @@ func TestJSONStoreSaveReplaceFailureWhenTargetIsNonEmptyDirectory(t *testing.T) CreatedAt: time.Now(), UpdatedAt: time.Now(), }) - if err == nil || !strings.Contains(err.Error(), "replace session file") { + if err == nil || !strings.Contains(err.Error(), "replace file") { t.Fatalf("expected replace failure, got %v", err) } } @@ -334,18 +337,23 @@ func TestJSONStoreSaveOverwritesExistingSessionFile(t *testing.T) { func TestJSONStoreSaveWriteTempFailure(t *testing.T) { t.Parallel() + if goruntime.GOOS == "windows" { + t.Skip("chmod write permission behavior is platform-specific on windows") + } baseDir := t.TempDir() workspaceRoot := t.TempDir() store := NewJSONStore(baseDir, workspaceRoot) - sessionsPath := sessionDirectory(baseDir, workspaceRoot) - if err := os.MkdirAll(sessionsPath, 0o755); err != nil { - t.Fatalf("mkdir sessions path: %v", err) + sessionDir := filepath.Join(sessionDirectory(baseDir, workspaceRoot), "temp-blocked") + if err := os.MkdirAll(sessionDir, 0o755); err != nil { + t.Fatalf("mkdir session dir: %v", err) } - tempDir := filepath.Join(sessionsPath, "temp-blocked.json.tmp") - if err := os.MkdirAll(tempDir, 0o755); err != nil { - t.Fatalf("mkdir temp dir: %v", err) + if err := os.Chmod(sessionDir, 0o555); err != nil { + t.Fatalf("chmod session dir readonly: %v", err) } + t.Cleanup(func() { + _ = os.Chmod(sessionDir, 0o755) + }) err := store.Save(context.Background(), &Session{ SchemaVersion: CurrentSchemaVersion, @@ -354,11 +362,22 @@ func TestJSONStoreSaveWriteTempFailure(t *testing.T) { CreatedAt: time.Now(), UpdatedAt: time.Now(), }) - if err == nil || !strings.Contains(err.Error(), "write temp session") { + if err == nil || !strings.Contains(err.Error(), "create temp file") { t.Fatalf("expected temp write failure, got %v", err) } } +func TestEnsurePathWithinBaseRejectsEscapedPath(t *testing.T) { + t.Parallel() + + baseDir := t.TempDir() + escaped := filepath.Join(baseDir, "..", "escaped", "session.json") + + if err := ensurePathWithinBase(baseDir, escaped); err == nil || !strings.Contains(err.Error(), "escapes base dir") { + t.Fatalf("expected escaped path rejection, got %v", err) + } +} + func TestJSONStoreLoadMissingFileReturnsError(t *testing.T) { t.Parallel() @@ -368,6 +387,15 @@ func TestJSONStoreLoadMissingFileReturnsError(t *testing.T) { } } +func TestJSONStoreLoadRejectsInvalidSessionID(t *testing.T) { + t.Parallel() + + store := NewJSONStore(t.TempDir(), t.TempDir()) + if _, err := store.Load(context.Background(), "bad/id"); err == nil || !strings.Contains(err.Error(), "unsupported characters") { + t.Fatalf("expected invalid session id error, got %v", err) + } +} + func TestNewUsesDefaultWorkdirAndEmptyMessages(t *testing.T) { t.Parallel() @@ -488,7 +516,7 @@ func TestJSONStoreLoadDecodeErrorWithNonJSONPayload(t *testing.T) { workspaceRoot := t.TempDir() store := NewJSONStore(baseDir, workspaceRoot) - mustWriteSessionFile(t, filepath.Join(sessionDirectory(baseDir, workspaceRoot), "decode-bad.json"), "{not-json") + mustWriteSessionFile(t, sessionFilePathForTest(baseDir, workspaceRoot, "decode-bad"), "{not-json") _, err := store.Load(context.Background(), "decode-bad") if err == nil || !strings.Contains(err.Error(), "decode session decode-bad") { @@ -505,7 +533,7 @@ func TestJSONStoreLoadRejectsMissingSchemaVersion(t *testing.T) { mustWriteSessionFile( t, - filepath.Join(sessionDirectory(baseDir, workspaceRoot), "missing-schema.json"), + sessionFilePathForTest(baseDir, workspaceRoot, "missing-schema"), `{"id":"missing-schema","title":"x","task_state":{"goal":"","progress":[],"open_items":[],"next_step":"","blockers":[],"key_artifacts":[],"decisions":[],"user_constraints":[],"last_updated_at":"0001-01-01T00:00:00Z"},"messages":[]}`, ) @@ -524,7 +552,7 @@ func TestJSONStoreLoadRejectsMissingTaskState(t *testing.T) { mustWriteSessionFile( t, - filepath.Join(sessionDirectory(baseDir, workspaceRoot), "missing-task-state.json"), + sessionFilePathForTest(baseDir, workspaceRoot, "missing-task-state"), `{"schema_version":2,"id":"missing-task-state","title":"x","messages":[]}`, ) @@ -552,11 +580,15 @@ func TestJSONStoreListSummariesSkipsUnreadableAndMalformedEntries(t *testing.T) t.Fatalf("save valid session: %v", err) } - mustWriteSessionFile(t, filepath.Join(sessionDirectory(baseDir, workspaceRoot), "malformed.json"), "{malformed") - mustWriteSessionFile(t, filepath.Join(sessionDirectory(baseDir, workspaceRoot), "empty-id.json"), `{"id":" ","title":"x"}`) + mustWriteSessionFile(t, sessionFilePathForTest(baseDir, workspaceRoot, "malformed"), "{malformed") + mustWriteSessionFile( + t, + sessionFilePathForTest(baseDir, workspaceRoot, "empty-id"), + `{"schema_version":2,"id":" ","title":"x","created_at":"2026-04-13T00:00:00Z","updated_at":"2026-04-13T00:00:00Z","task_state":{"goal":"","progress":[],"open_items":[],"next_step":"","blockers":[],"key_artifacts":[],"decisions":[],"user_constraints":[],"last_updated_at":"2026-04-13T00:00:00Z"}}`, + ) mustWriteSessionFile( t, - filepath.Join(sessionDirectory(baseDir, workspaceRoot), "missing-task-state-summary.json"), + sessionFilePathForTest(baseDir, workspaceRoot, "missing-task-state-summary"), `{"schema_version":2,"id":"missing-task-state-summary","title":"x","created_at":"2026-04-13T00:00:00Z","updated_at":"2026-04-13T00:00:00Z"}`, ) @@ -569,6 +601,111 @@ func TestJSONStoreListSummariesSkipsUnreadableAndMalformedEntries(t *testing.T) } } +func TestJSONStoreSaveRejectsInvalidSchemaAndID(t *testing.T) { + t.Parallel() + + store := NewJSONStore(t.TempDir(), t.TempDir()) + + err := store.Save(context.Background(), &Session{ + SchemaVersion: CurrentSchemaVersion + 1, + ID: "session-invalid-schema", + Title: "Invalid Schema", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + if err == nil || !strings.Contains(err.Error(), "unsupported schema_version") { + t.Fatalf("expected schema version error, got %v", err) + } + + err = store.Save(context.Background(), &Session{ + SchemaVersion: CurrentSchemaVersion, + ID: "bad/session", + Title: "Invalid ID", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + if err == nil || !strings.Contains(err.Error(), "unsupported characters") { + t.Fatalf("expected invalid id error, got %v", err) + } +} + +func TestJSONStoreSaveCreateSessionDirFailure(t *testing.T) { + t.Parallel() + + baseDir := t.TempDir() + workspaceRoot := t.TempDir() + store := NewJSONStore(baseDir, workspaceRoot) + sessionID := "session-dir-failed" + + sessionDirPath := filepath.Join(sessionDirectory(baseDir, workspaceRoot), sessionID) + if err := os.MkdirAll(filepath.Dir(sessionDirPath), 0o755); err != nil { + t.Fatalf("mkdir session parent: %v", err) + } + if err := os.WriteFile(sessionDirPath, []byte("blocked"), 0o644); err != nil { + t.Fatalf("write session dir blocker: %v", err) + } + + err := store.Save(context.Background(), &Session{ + SchemaVersion: CurrentSchemaVersion, + ID: sessionID, + Title: "Blocked SessionDir", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + if err == nil || !strings.Contains(err.Error(), "create session dir") { + t.Fatalf("expected create session dir error, got %v", err) + } +} + +func TestDecodeStoredSessionAndSummarySchemaValidation(t *testing.T) { + t.Parallel() + + _, err := decodeStoredSession([]byte(`{ + "schema_version": 999, + "id": "decode-invalid-schema", + "title": "x", + "created_at": "2026-04-13T08:00:00Z", + "updated_at": "2026-04-13T08:00:00Z", + "task_state": { + "goal": "", + "progress": [], + "open_items": [], + "next_step": "", + "blockers": [], + "key_artifacts": [], + "decisions": [], + "user_constraints": [], + "last_updated_at": "2026-04-13T08:00:00Z" + }, + "messages": [] +}`)) + if err == nil || !strings.Contains(err.Error(), "unsupported schema_version") { + t.Fatalf("expected decodeStoredSession schema error, got %v", err) + } + + _, err = decodeStoredSummary([]byte(`{ + "schema_version": 999, + "id": "summary-invalid-schema", + "title": "Summary", + "created_at": "2026-04-13T08:00:00Z", + "updated_at": "2026-04-13T08:00:00Z", + "task_state": { + "goal": "", + "progress": [], + "open_items": [], + "next_step": "", + "blockers": [], + "key_artifacts": [], + "decisions": [], + "user_constraints": [], + "last_updated_at": "2026-04-13T08:00:00Z" + } +}`)) + if err == nil || !strings.Contains(err.Error(), "unsupported schema_version") { + t.Fatalf("expected decodeStoredSummary schema error, got %v", err) + } +} + func TestJSONStoreSavePersistsProviderModelAndMessages(t *testing.T) { t.Parallel() @@ -610,7 +747,7 @@ func TestJSONStoreSavePersistsProviderModelAndMessages(t *testing.T) { t.Fatalf("save session: %v", err) } - rawPath := filepath.Join(sessionDirectory(baseDir, workspaceRoot), session.ID+".json") + rawPath := sessionFilePathForTest(baseDir, workspaceRoot, session.ID) raw, err := os.ReadFile(rawPath) if err != nil { t.Fatalf("read raw file: %v", err) @@ -803,7 +940,7 @@ func TestJSONStoreLoadClampsOversizedTaskState(t *testing.T) { }, "\n") mustWriteSessionFile( t, - filepath.Join(sessionDirectory(baseDir, workspaceRoot), "task-state-clamp-load.json"), + sessionFilePathForTest(baseDir, workspaceRoot, "task-state-clamp-load"), payload, ) @@ -887,7 +1024,7 @@ func TestJSONStoreLoadAllowsMissingTodosField(t *testing.T) { workspaceRoot := t.TempDir() store := NewJSONStore(baseDir, workspaceRoot) - mustWriteSessionFile(t, filepath.Join(sessionDirectory(baseDir, workspaceRoot), "no-todos.json"), strings.Join([]string{ + mustWriteSessionFile(t, sessionFilePathForTest(baseDir, workspaceRoot, "no-todos"), strings.Join([]string{ `{`, ` "schema_version": 2,`, ` "id": "no-todos",`, @@ -918,6 +1055,53 @@ func TestJSONStoreLoadAllowsMissingTodosField(t *testing.T) { } } +func TestJSONStoreLoadBackfillsTodoVersionWhenMissing(t *testing.T) { + t.Parallel() + + baseDir := t.TempDir() + workspaceRoot := t.TempDir() + store := NewJSONStore(baseDir, workspaceRoot) + + mustWriteSessionFile(t, sessionFilePathForTest(baseDir, workspaceRoot, "todos-no-version"), strings.Join([]string{ + `{`, + ` "schema_version": 2,`, + ` "id": "todos-no-version",`, + ` "title": "Todos No Version",`, + ` "created_at": "2026-04-14T10:00:00Z",`, + ` "updated_at": "2026-04-14T10:05:00Z",`, + ` "task_state": {`, + ` "goal": "",`, + ` "progress": [],`, + ` "open_items": [],`, + ` "next_step": "",`, + ` "blockers": [],`, + ` "key_artifacts": [],`, + ` "decisions": [],`, + ` "user_constraints": [],`, + ` "last_updated_at": "2026-04-14T10:05:00Z"`, + ` },`, + ` "todos": [`, + ` {`, + ` "id": "todo-1",`, + ` "content": "todo item",`, + ` "status": "pending",`, + ` "created_at": "2026-04-14T10:00:00Z",`, + ` "updated_at": "2026-04-14T10:05:00Z"`, + ` }`, + ` ],`, + ` "messages": []`, + `}`, + }, "\n")) + + loaded, err := store.Load(context.Background(), "todos-no-version") + if err != nil { + t.Fatalf("load session with todos and missing todo_version: %v", err) + } + if loaded.TodoVersion != CurrentTodoVersion { + t.Fatalf("expected todo version %d, got %d", CurrentTodoVersion, loaded.TodoVersion) + } +} + func TestJSONStoreSaveRejectsInvalidTodos(t *testing.T) { t.Parallel() @@ -941,6 +1125,76 @@ func TestJSONStoreSaveRejectsInvalidTodos(t *testing.T) { } } +func TestDecodeStoredSummaryRejectsMissingSchemaVersion(t *testing.T) { + t.Parallel() + + _, err := decodeStoredSummary([]byte(`{"id":"summary-no-schema","task_state":{}}`)) + if err == nil || !strings.Contains(err.Error(), "missing required field schema_version") { + t.Fatalf("expected missing schema_version error, got %v", err) + } +} + +func TestCreateTempFileAndAtomicReplaceFailureBranches(t *testing.T) { + t.Parallel() + + baseDir := t.TempDir() + missingDir := filepath.Join(baseDir, "missing", "dir") + if _, _, err := createTempFile(missingDir, "tmp-*.tmp", "create temp file"); err == nil || + !strings.Contains(err.Error(), "create temp file") { + t.Fatalf("expected create temp file error for missing dir, got %v", err) + } + + source := filepath.Join(baseDir, "source.tmp") + if err := os.WriteFile(source, []byte("payload"), 0o644); err != nil { + t.Fatalf("write source temp: %v", err) + } + blockerDir := filepath.Join(baseDir, "target") + if err := os.MkdirAll(filepath.Join(blockerDir, "child"), 0o755); err != nil { + t.Fatalf("mkdir blocker dir: %v", err) + } + if err := replaceFileWithTemp(source, blockerDir, "file"); err == nil || !strings.Contains(err.Error(), "replace file") { + t.Fatalf("expected replace failure for non-empty target dir, got %v", err) + } +} + +func TestReplaceFileWithTempCommitFailure(t *testing.T) { + t.Parallel() + + baseDir := t.TempDir() + target := filepath.Join(baseDir, "target.txt") + if err := replaceFileWithTemp(filepath.Join(baseDir, "missing.tmp"), target, "file"); err == nil || + !strings.Contains(err.Error(), "commit file") { + t.Fatalf("expected commit failure when temp file missing, got %v", err) + } +} + +func TestWriteFileAtomicallyCreateTempFailure(t *testing.T) { + t.Parallel() + + target := filepath.Join(t.TempDir(), "missing", "nested", "session.json") + err := writeFileAtomically(target, "session-*.tmp", []byte("data"), 0o644) + if err == nil || !strings.Contains(err.Error(), "create temp file") { + t.Fatalf("expected create temp file error, got %v", err) + } +} + +func TestJSONStoreRejectsInvalidBasePathDuringAssetAndSessionResolve(t *testing.T) { + t.Parallel() + + store := &JSONStore{baseDir: string([]byte{'b', 'a', 'd', 0})} + if _, err := store.Load(context.Background(), "session-ok"); err == nil { + t.Fatalf("expected load error for invalid base path") + } + + if _, err := store.SaveAsset(context.Background(), "session-ok", strings.NewReader("img"), "image/png"); err == nil { + t.Fatalf("expected save asset error for invalid base path") + } + + if _, err := store.Stat(context.Background(), "session-ok", "asset-ok"); err == nil { + t.Fatalf("expected stat error for invalid base path") + } +} + func TestJSONStoreLoadRejectsInvalidTodos(t *testing.T) { t.Parallel() @@ -948,7 +1202,7 @@ func TestJSONStoreLoadRejectsInvalidTodos(t *testing.T) { workspaceRoot := t.TempDir() store := NewJSONStore(baseDir, workspaceRoot) - mustWriteSessionFile(t, filepath.Join(sessionDirectory(baseDir, workspaceRoot), "invalid-todos-load.json"), strings.Join([]string{ + mustWriteSessionFile(t, sessionFilePathForTest(baseDir, workspaceRoot, "invalid-todos-load"), strings.Join([]string{ `{`, ` "schema_version": 2,`, ` "id": "invalid-todos-load",`, @@ -1009,3 +1263,7 @@ func mustWriteSessionFile(t *testing.T, path string, content string) { t.Fatalf("write %s: %v", path, err) } } + +func sessionFilePathForTest(baseDir string, workspaceRoot string, sessionID string) string { + return filepath.Join(sessionDirectory(baseDir, workspaceRoot), sessionID, sessionFileName) +}