211 lines
9.2 KiB
Go
211 lines
9.2 KiB
Go
package models
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"testing"
|
|
|
|
"git.cin.su/m/Assistant/contracts/authority"
|
|
"git.cin.su/m/Assistant/contracts/capability"
|
|
"git.cin.su/m/Assistant/contracts/cognitive"
|
|
"git.cin.su/m/Assistant/contracts/identity"
|
|
)
|
|
|
|
type roundTripFunc func(*http.Request) (*http.Response, error)
|
|
|
|
func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { return f(request) }
|
|
|
|
func validOpenAIAdvisoryRequest() cognitive.AdvisoryRequest {
|
|
return cognitive.AdvisoryRequest{
|
|
Version: cognitive.CreatorAdvisoryV1,
|
|
RequestID: "request-1",
|
|
ActorID: identity.ActorID("actor-private"),
|
|
PrincipalID: identity.PrincipalID("principal-private"),
|
|
ScopeID: authority.ScopeID("scope-private"),
|
|
Message: "What should I verify?",
|
|
Locale: "en",
|
|
MaxOutputBytes: MaxAdvisoryResponseBytes,
|
|
ProviderBindingID: "binding-under-test",
|
|
ProviderBindingRevision: 7,
|
|
ProviderInstanceID: capability.ProviderInstanceID("provider-under-test"),
|
|
ModelDescriptorID: capability.ModelDescriptorID("model-under-test"),
|
|
ProviderModelID: capability.ProviderModelID(openAIModel),
|
|
}
|
|
}
|
|
|
|
func providerHTTPResponse(status int, body string) *http.Response {
|
|
return &http.Response{
|
|
StatusCode: status,
|
|
Header: http.Header{"Content-Type": []string{"application/json"}},
|
|
Body: io.NopCloser(strings.NewReader(body)),
|
|
}
|
|
}
|
|
|
|
func TestOpenAIStatusTaxonomyPreservesRateLimit(t *testing.T) {
|
|
if got := mapOpenAIStatus(http.StatusTooManyRequests); got != cognitive.FinishRateLimited {
|
|
t.Fatalf("status 429 mapped to %s", got)
|
|
}
|
|
}
|
|
|
|
func validOpenAIResponse(text string) string {
|
|
body, _ := json.Marshal(map[string]any{
|
|
"status": "completed",
|
|
"model": openAIModel,
|
|
"store": false,
|
|
"output": []any{
|
|
map[string]any{"type": "reasoning", "summary": []any{}},
|
|
map[string]any{"type": "message", "role": "assistant", "status": "completed", "content": []any{
|
|
map[string]any{"type": "output_text", "text": text, "annotations": []any{}},
|
|
}},
|
|
},
|
|
"usage": map[string]any{"input_tokens": 12, "output_tokens": 7},
|
|
})
|
|
return string(body)
|
|
}
|
|
|
|
func TestOpenAIAdvisorySendsOnlyModeAPayload(t *testing.T) {
|
|
calls := 0
|
|
transport := roundTripFunc(func(request *http.Request) (*http.Response, error) {
|
|
calls++
|
|
if request.Method != http.MethodPost || request.URL.String() != openAIResponsesURL || request.Host != "api.openai.com" {
|
|
t.Fatalf("unexpected target: %s %s host=%s", request.Method, request.URL, request.Host)
|
|
}
|
|
if request.Header.Get("Authorization") != "Bearer test-key" || request.Header.Get("Content-Type") != "application/json" {
|
|
t.Fatal("provider authentication or content type is missing")
|
|
}
|
|
body, err := io.ReadAll(request.Body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var payload map[string]any
|
|
if err := json.Unmarshal(body, &payload); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(payload) != 6 || payload["model"] != openAIModel || payload["input"] != "What should I verify?" || payload["store"] != false || payload["instructions"] != openAIInstructions || payload["max_output_tokens"] != float64(openAIMaxOutputTokens) {
|
|
t.Fatalf("unexpected provider payload keys/values: %v", payload)
|
|
}
|
|
reasoning, ok := payload["reasoning"].(map[string]any)
|
|
if !ok || len(reasoning) != 1 || reasoning["effort"] != "low" {
|
|
t.Fatalf("unexpected reasoning payload: %v", payload["reasoning"])
|
|
}
|
|
serialized := string(body)
|
|
for _, forbidden := range []string{"actor-private", "principal-private", "scope-private", "request-1", "tools", "previous_response_id", "conversation", "background"} {
|
|
if strings.Contains(serialized, forbidden) {
|
|
t.Fatalf("provider payload leaked forbidden field %q", forbidden)
|
|
}
|
|
}
|
|
return providerHTTPResponse(http.StatusOK, validOpenAIResponse("Check the current evidence.")), nil
|
|
})
|
|
adapter, err := NewOpenAIAdvisoryAdapter([]byte("test-key"), transport)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
response, err := adapter.Advise(context.Background(), validOpenAIAdvisoryRequest())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if calls != 1 || response.FinishReason != cognitive.FinishCompleted || response.AdvisoryText != "Check the current evidence." || response.Usage == nil || response.Usage.InputUnits != 12 || response.Usage.OutputUnits != 7 {
|
|
t.Fatalf("unexpected response: calls=%d response=%+v", calls, response)
|
|
}
|
|
}
|
|
|
|
func TestOpenAIAdvisoryRejectsContextBeforeTransport(t *testing.T) {
|
|
called := false
|
|
adapter, _ := NewOpenAIAdvisoryAdapter([]byte("test-key"), roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
called = true
|
|
return nil, errors.New("must not call")
|
|
}))
|
|
request := validOpenAIAdvisoryRequest()
|
|
request.Context = []cognitive.ContextExcerpt{{ID: "excerpt", SourceKind: "task", SourceID: "task", Revision: 1, Sensitivity: cognitive.SensitivityPublic, Text: "public"}}
|
|
if _, err := adapter.Advise(context.Background(), request); err == nil || called {
|
|
t.Fatalf("error=%v called=%t", err, called)
|
|
}
|
|
}
|
|
|
|
func TestOpenAIAdvisoryDoesNotFollowRedirectOrRetry(t *testing.T) {
|
|
calls := 0
|
|
adapter, _ := NewOpenAIAdvisoryAdapter([]byte("test-key"), roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
calls++
|
|
response := providerHTTPResponse(http.StatusFound, "{}")
|
|
response.Header.Set("Location", "https://attacker.invalid/steal")
|
|
return response, nil
|
|
}))
|
|
response, err := adapter.Advise(context.Background(), validOpenAIAdvisoryRequest())
|
|
if err != nil || calls != 1 || response.FinishReason != cognitive.FinishProviderUnavailable {
|
|
t.Fatalf("error=%v calls=%d response=%+v", err, calls, response)
|
|
}
|
|
}
|
|
|
|
type timeoutError struct{}
|
|
|
|
func (timeoutError) Error() string { return "timeout" }
|
|
func (timeoutError) Timeout() bool { return true }
|
|
func (timeoutError) Temporary() bool { return true }
|
|
|
|
func TestOpenAIAdvisoryMapsTimeoutAndDoesNotRetry(t *testing.T) {
|
|
calls := 0
|
|
adapter, _ := NewOpenAIAdvisoryAdapter([]byte("test-key"), roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
calls++
|
|
return nil, timeoutError{}
|
|
}))
|
|
response, err := adapter.Advise(context.Background(), validOpenAIAdvisoryRequest())
|
|
if err != nil || calls != 1 || response.FinishReason != cognitive.FinishTimeout {
|
|
t.Fatalf("error=%v calls=%d response=%+v", err, calls, response)
|
|
}
|
|
}
|
|
|
|
func TestOpenAIAdvisoryMapsRefusal(t *testing.T) {
|
|
body := `{"status":"completed","model":"gpt-5.4-mini-2026-03-17","store":false,"output":[{"type":"message","role":"assistant","status":"completed","content":[{"type":"refusal","refusal":"cannot comply"}]}],"usage":{"input_tokens":3,"output_tokens":0}}`
|
|
adapter, _ := NewOpenAIAdvisoryAdapter([]byte("test-key"), roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
return providerHTTPResponse(http.StatusOK, body), nil
|
|
}))
|
|
response, err := adapter.Advise(context.Background(), validOpenAIAdvisoryRequest())
|
|
if err != nil || response.FinishReason != cognitive.FinishContentRejected || response.AdvisoryText != "" {
|
|
t.Fatalf("error=%v response=%+v", err, response)
|
|
}
|
|
}
|
|
|
|
func TestOpenAIAdvisoryMapsOutputLimit(t *testing.T) {
|
|
body := `{"status":"incomplete","model":"gpt-5.4-mini-2026-03-17","store":false,"incomplete_details":{"reason":"max_output_tokens"},"output":[],"usage":{"input_tokens":3,"output_tokens":4096}}`
|
|
adapter, _ := NewOpenAIAdvisoryAdapter([]byte("test-key"), roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
return providerHTTPResponse(http.StatusOK, body), nil
|
|
}))
|
|
response, err := adapter.Advise(context.Background(), validOpenAIAdvisoryRequest())
|
|
if err != nil || response.FinishReason != cognitive.FinishLengthLimit || response.AdvisoryText != "" {
|
|
t.Fatalf("error=%v response=%+v", err, response)
|
|
}
|
|
}
|
|
|
|
func TestOpenAIAdvisoryRejectsExecutableOrUntrustedOutput(t *testing.T) {
|
|
tests := map[string]string{
|
|
"tool call": `{"status":"completed","model":"gpt-5.4-mini-2026-03-17","store":false,"output":[{"type":"function_call","name":"mutate"}],"usage":{"input_tokens":1,"output_tokens":1}}`,
|
|
"citation": `{"status":"completed","model":"gpt-5.4-mini-2026-03-17","store":false,"output":[{"type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"answer","annotations":[{"type":"url_citation"}]}]}],"usage":{"input_tokens":1,"output_tokens":1}}`,
|
|
"wrong model": strings.Replace(validOpenAIResponse("answer"), openAIModel, "moving-alias", 1),
|
|
"oversize": validOpenAIResponse(strings.Repeat("a", MaxAdvisoryResponseBytes+1)),
|
|
"malformed": `{`,
|
|
}
|
|
for name, body := range tests {
|
|
t.Run(name, func(t *testing.T) {
|
|
adapter, _ := NewOpenAIAdvisoryAdapter([]byte("test-key"), roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
return providerHTTPResponse(http.StatusOK, body), nil
|
|
}))
|
|
response, err := adapter.Advise(context.Background(), validOpenAIAdvisoryRequest())
|
|
if err != nil || response.FinishReason != cognitive.FinishInvalidProviderResponse || response.AdvisoryText != "" {
|
|
t.Fatalf("error=%v response=%+v", err, response)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestOpenAIAdvisoryRejectsInvalidCredential(t *testing.T) {
|
|
for _, key := range [][]byte{nil, []byte("line\nbreak")} {
|
|
if _, err := NewOpenAIAdvisoryAdapter(key, nil); err == nil {
|
|
t.Fatalf("credential %q was accepted", key)
|
|
}
|
|
}
|
|
}
|