fix: validate Fabric chat URL host and tidy Ollama responses

## CHANGES
- Set NDJSON content type before checking upstream errors
- Reject parsed URLs that omit a hostname
- Remove hardcoded eval count placeholders from responses
- Add unit tests for Fabric chat URL builder
- Cover colon-port, host:port, and IP address inputs
This commit is contained in:
Kayvan Sylvan 2026-01-17 02:31:46 -08:00
parent ae6d4d1fb3
commit 29a32a8439
2 changed files with 95 additions and 8 deletions

View file

@ -217,6 +217,10 @@ func (f APIConvert) ollamaChat(c *gin.Context) {
}
defer fabricRes.Body.Close()
if prompt.Stream {
c.Header("Content-Type", "application/x-ndjson")
}
if fabricRes.StatusCode < http.StatusOK || fabricRes.StatusCode >= http.StatusMultipleChoices {
bodyBytes, readErr := io.ReadAll(fabricRes.Body)
if readErr != nil {
@ -234,10 +238,6 @@ func (f APIConvert) ollamaChat(c *gin.Context) {
return
}
if prompt.Stream {
c.Header("Content-Type", "application/x-ndjson")
}
var contentBuilder strings.Builder
scanner := bufio.NewScanner(fabricRes.Body)
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
@ -304,9 +304,7 @@ func (f APIConvert) ollamaChat(c *gin.Context) {
Done: true,
TotalDuration: time.Since(now).Nanoseconds(),
LoadDuration: int(time.Since(now).Nanoseconds()),
PromptEvalCount: 42,
PromptEvalDuration: int(time.Since(now).Nanoseconds()),
EvalCount: 420,
EvalDuration: time.Since(now).Nanoseconds(),
}
c.JSON(200, response)
@ -327,9 +325,7 @@ func (f APIConvert) ollamaChat(c *gin.Context) {
Done: true,
TotalDuration: time.Since(now).Nanoseconds(),
LoadDuration: int(time.Since(now).Nanoseconds()),
PromptEvalCount: 42,
PromptEvalDuration: int(time.Since(now).Nanoseconds()),
EvalCount: 420,
EvalDuration: time.Since(now).Nanoseconds(),
}
if err := writeOllamaResponseStruct(c, finalResponse); err != nil {
@ -346,6 +342,9 @@ func buildFabricChatURL(addr string) (string, error) {
if err != nil {
return "", fmt.Errorf("invalid address: %w", err)
}
if parsed.Host == "" {
return "", fmt.Errorf("invalid address: missing host")
}
return strings.TrimRight(parsed.String(), "/"), nil
}
if strings.HasPrefix(addr, ":") {

View file

@ -0,0 +1,88 @@
package restapi
import (
"testing"
)
func TestBuildFabricChatURL(t *testing.T) {
tests := []struct {
name string
addr string
want string
wantErr bool
}{
{
name: "empty address",
addr: "",
want: "",
wantErr: true,
},
{
name: "valid http URL",
addr: "http://localhost:8080",
want: "http://localhost:8080",
wantErr: false,
},
{
name: "valid https URL",
addr: "https://api.example.com",
want: "https://api.example.com",
wantErr: false,
},
{
name: "http URL with trailing slash",
addr: "http://localhost:8080/",
want: "http://localhost:8080",
wantErr: false,
},
{
name: "malformed URL - missing host",
addr: "http://",
want: "",
wantErr: true,
},
{
name: "malformed URL - port only with http",
addr: "https://:8080",
want: "https://:8080",
wantErr: false,
},
{
name: "colon-prefixed port",
addr: ":8080",
want: "http://127.0.0.1:8080",
wantErr: false,
},
{
name: "bare host:port",
addr: "localhost:8080",
want: "http://localhost:8080",
wantErr: false,
},
{
name: "bare hostname",
addr: "localhost",
want: "http://localhost",
wantErr: false,
},
{
name: "IP address with port",
addr: "192.168.1.1:3000",
want: "http://192.168.1.1:3000",
wantErr: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := buildFabricChatURL(tt.addr)
if (err != nil) != tt.wantErr {
t.Errorf("buildFabricChatURL() error = %v, wantErr %v", err, tt.wantErr)
return
}
if got != tt.want {
t.Errorf("buildFabricChatURL() = %v, want %v", got, tt.want)
}
})
}
}