Skip to content

Commit 2aeb41c

Browse files
luispaterluispater
committed
feat(pluginhost, jshandler): integrate HostCallbackID with interceptors and JS engine logging
- Added `HostCallbackID` to request, response, and stream chunk interceptors for enhanced context tracking. - Updated JavaScript engine to support custom console logging with `HostCallbackID` forwarding. - Introduced tests verifying proper integration of `HostCallbackID` in all interceptor flows and engine logging. - Enhanced logging and error handling for consistent callback-related logic implementation.
1 parent 1d1ee85 commit 2aeb41c

9 files changed

Lines changed: 359 additions & 19 deletions

File tree

‎examples/plugin/jshandler/abi.go‎

Lines changed: 110 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,28 @@ type abiLifecycleRequest struct {
9292
PluginDir string `json:"plugin_dir,omitempty"`
9393
}
9494

95+
type abiRequestInterceptRequest struct {
96+
pluginapi.RequestInterceptRequest
97+
HostCallbackID string `json:"host_callback_id,omitempty"`
98+
}
99+
100+
type abiResponseInterceptRequest struct {
101+
pluginapi.ResponseInterceptRequest
102+
HostCallbackID string `json:"host_callback_id,omitempty"`
103+
}
104+
105+
type abiStreamChunkInterceptRequest struct {
106+
pluginapi.StreamChunkInterceptRequest
107+
HostCallbackID string `json:"host_callback_id,omitempty"`
108+
}
109+
110+
type abiHostLogRequest struct {
111+
HostCallbackID string `json:"host_callback_id,omitempty"`
112+
Level string `json:"level,omitempty"`
113+
Message string `json:"message,omitempty"`
114+
Fields map[string]any `json:"fields,omitempty"`
115+
}
116+
95117
type abiRegistration struct {
96118
SchemaVersion uint32 `json:"schema_version"`
97119
Metadata pluginapi.Metadata `json:"metadata"`
@@ -192,25 +214,25 @@ func handleJSHandlerABIMethod(ctx context.Context, method string, request []byte
192214
defer done()
193215
switch method {
194216
case pluginabi.MethodRequestInterceptBefore:
195-
var req pluginapi.RequestInterceptRequest
217+
var req abiRequestInterceptRequest
196218
if errDecode := json.Unmarshal(request, &req); errDecode != nil {
197219
return nil, errDecode
198220
}
199-
resp, errCall := p.InterceptRequest(ctx, req)
221+
resp, errCall := p.interceptRequest(ctx, req.RequestInterceptRequest, req.HostCallbackID)
200222
return abiOKEnvelopeWithError(resp, errCall)
201223
case pluginabi.MethodResponseInterceptAfter:
202-
var req pluginapi.ResponseInterceptRequest
224+
var req abiResponseInterceptRequest
203225
if errDecode := json.Unmarshal(request, &req); errDecode != nil {
204226
return nil, errDecode
205227
}
206-
resp, errCall := p.InterceptResponse(ctx, req)
228+
resp, errCall := p.interceptResponse(ctx, req.ResponseInterceptRequest, req.HostCallbackID)
207229
return abiOKEnvelopeWithError(resp, errCall)
208230
case pluginabi.MethodResponseInterceptStreamChunk:
209-
var req pluginapi.StreamChunkInterceptRequest
231+
var req abiStreamChunkInterceptRequest
210232
if errDecode := json.Unmarshal(request, &req); errDecode != nil {
211233
return nil, errDecode
212234
}
213-
resp, errCall := p.InterceptStreamChunk(ctx, req)
235+
resp, errCall := p.interceptStreamChunk(ctx, req.StreamChunkInterceptRequest, req.HostCallbackID)
214236
return abiOKEnvelopeWithError(resp, errCall)
215237
default:
216238
return abiErrorEnvelope("unknown_method", "unknown method: "+method), nil
@@ -289,3 +311,85 @@ func writeABIResponse(response *C.cliproxy_buffer, raw []byte) {
289311
response.ptr = ptr
290312
response.len = C.size_t(len(raw))
291313
}
314+
315+
func newHostJSConsoleLogger(hostCallbackID string) jsConsoleLogger {
316+
return func(message string) error {
317+
if errLog := writeHostJSConsoleLog(hostCallbackID, message); errLog != nil {
318+
return defaultJSConsoleLogger(message)
319+
}
320+
return nil
321+
}
322+
}
323+
324+
func writeHostJSConsoleLog(hostCallbackID string, message string) error {
325+
raw, errMarshal := json.Marshal(abiHostLogRequest{
326+
HostCallbackID: hostCallbackID,
327+
Level: "info",
328+
Message: "JS console log: " + message,
329+
Fields: map[string]any{
330+
"plugin_id": pluginName,
331+
},
332+
})
333+
if errMarshal != nil {
334+
return errMarshal
335+
}
336+
337+
rawResp, errCall := callHost(pluginabi.MethodHostLog, raw)
338+
if errCall != nil {
339+
return errCall
340+
}
341+
if len(rawResp) == 0 {
342+
return nil
343+
}
344+
var resp abiEnvelope
345+
if errDecode := json.Unmarshal(rawResp, &resp); errDecode != nil {
346+
return fmt.Errorf("decode host log response: %w", errDecode)
347+
}
348+
if !resp.OK {
349+
if resp.Error != nil {
350+
return fmt.Errorf("host log failed: %s", resp.Error.Message)
351+
}
352+
return fmt.Errorf("host log failed")
353+
}
354+
return nil
355+
}
356+
357+
func callHost(method string, payload []byte) ([]byte, error) {
358+
jsHandlerABIState.RLock()
359+
defer jsHandlerABIState.RUnlock()
360+
if jsHandlerABIState.host == nil {
361+
return nil, fmt.Errorf("host callback is unavailable")
362+
}
363+
364+
cMethod := C.CString(method)
365+
defer C.free(unsafe.Pointer(cMethod))
366+
367+
var cPayload unsafe.Pointer
368+
if len(payload) > 0 {
369+
cPayload = C.CBytes(payload)
370+
if cPayload == nil {
371+
return nil, fmt.Errorf("allocate host callback payload")
372+
}
373+
defer C.free(cPayload)
374+
}
375+
376+
var response C.cliproxy_buffer
377+
rc := C.jshandler_call_host(
378+
jsHandlerABIState.host,
379+
cMethod,
380+
(*C.uint8_t)(cPayload),
381+
C.size_t(len(payload)),
382+
&response,
383+
)
384+
var out []byte
385+
if response.ptr != nil && response.len > 0 {
386+
out = C.GoBytes(response.ptr, C.int(response.len))
387+
}
388+
if response.ptr != nil {
389+
C.jshandler_free_host_buffer(jsHandlerABIState.host, response.ptr, response.len)
390+
}
391+
if rc != 0 {
392+
return nil, fmt.Errorf("host callback %s returned %d: %s", method, int(rc), string(out))
393+
}
394+
return out, nil
395+
}

‎examples/plugin/jshandler/engine.go‎

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"fmt"
66
"os"
77
"path/filepath"
8+
"strings"
89
"sync"
910
"time"
1011

@@ -13,27 +14,43 @@ import (
1314
)
1415

1516
type jsEngine struct {
16-
vm *goja.Runtime
17+
vm *goja.Runtime
18+
consoleLogger jsConsoleLogger
1719
}
1820

1921
const maxJSScriptBytes = 8 * 1024 * 1024
2022

21-
func newJSEngine() *jsEngine {
23+
type jsConsoleLogger func(message string) error
24+
25+
func newJSEngine(loggers ...jsConsoleLogger) *jsEngine {
26+
consoleLogger := defaultJSConsoleLogger
27+
if len(loggers) > 0 && loggers[0] != nil {
28+
consoleLogger = loggers[0]
29+
}
2230
engine := &jsEngine{
23-
vm: goja.New(),
31+
vm: goja.New(),
32+
consoleLogger: consoleLogger,
2433
}
2534
engine.initConsole()
2635
return engine
2736
}
2837

38+
func defaultJSConsoleLogger(message string) error {
39+
log.Info("JS console log: ", message)
40+
return nil
41+
}
42+
2943
func (engine *jsEngine) initConsole() {
3044
console := engine.vm.NewObject()
3145
consoleLogWrapper := func(call goja.FunctionCall) goja.Value {
32-
args := make([]interface{}, len(call.Arguments))
46+
args := make([]string, len(call.Arguments))
3347
for i, arg := range call.Arguments {
34-
args[i] = arg.Export()
48+
args[i] = fmt.Sprint(arg.Export())
49+
}
50+
message := strings.Join(args, " ")
51+
if errLog := engine.consoleLogger(message); errLog != nil {
52+
defaultJSConsoleLogger(message)
3553
}
36-
log.Info("JS console log: ", fmt.Sprint(args...))
3754
return goja.Undefined()
3855
}
3956
_ = console.Set("log", consoleLogWrapper)

‎examples/plugin/jshandler/engine_test.go‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,59 @@
11
package main
22

33
import (
4+
"bytes"
5+
"strings"
46
"testing"
57
"time"
8+
9+
log "github.com/sirupsen/logrus"
610
)
711

12+
func TestConsoleLogWritesToLogger(t *testing.T) {
13+
var out bytes.Buffer
14+
logger := log.StandardLogger()
15+
originalOut := logger.Out
16+
originalFormatter := logger.Formatter
17+
originalLevel := logger.Level
18+
log.SetOutput(&out)
19+
log.SetFormatter(&log.TextFormatter{
20+
DisableColors: true,
21+
DisableTimestamp: true,
22+
})
23+
log.SetLevel(log.InfoLevel)
24+
defer func() {
25+
log.SetOutput(originalOut)
26+
log.SetFormatter(originalFormatter)
27+
log.SetLevel(originalLevel)
28+
}()
29+
30+
engine := newJSEngine()
31+
_, errRun := engine.vm.RunString(`console.log("alpha", 42, true);`)
32+
if errRun != nil {
33+
t.Fatalf("RunString() error = %v", errRun)
34+
}
35+
36+
got := out.String()
37+
if !strings.Contains(got, "JS console log: alpha 42 true") {
38+
t.Fatalf("console.log output = %q, want logger output with JS message", got)
39+
}
40+
}
41+
42+
func TestConsoleLogUsesConfiguredLogger(t *testing.T) {
43+
var messages []string
44+
engine := newJSEngine(func(message string) error {
45+
messages = append(messages, message)
46+
return nil
47+
})
48+
_, errRun := engine.vm.RunString(`console.log("alpha", 42, true);`)
49+
if errRun != nil {
50+
t.Fatalf("RunString() error = %v", errRun)
51+
}
52+
if len(messages) != 1 || messages[0] != "alpha 42 true" {
53+
t.Fatalf("console log messages = %#v, want formatted message", messages)
54+
}
55+
}
56+
857
func TestStopInterruptTimerClearsExpiredInterrupt(t *testing.T) {
958
engine := newJSEngine()
1059
timer, done := engine.startInterruptTimer(time.Nanosecond)

‎examples/plugin/jshandler/interceptor.go‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,10 @@ func (p *jsHandlerPlugin) allScriptPaths() []string {
4545
}
4646

4747
func (p *jsHandlerPlugin) InterceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) {
48+
return p.interceptRequest(ctx, req, "")
49+
}
50+
51+
func (p *jsHandlerPlugin) interceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest, hostCallbackID string) (pluginapi.RequestInterceptResponse, error) {
4852
resp := pluginapi.RequestInterceptResponse{}
4953
scriptPaths := p.allScriptPaths()
5054
if len(scriptPaths) == 0 {
@@ -60,7 +64,7 @@ func (p *jsHandlerPlugin) InterceptRequest(ctx context.Context, req pluginapi.Re
6064
if scriptPath == "" {
6165
continue
6266
}
63-
processed, cleared, errJS := p.applyJSBeforeRequest(scriptPath, []byte(body), req.Model, req.SourceFormat, headers)
67+
processed, cleared, errJS := p.applyJSBeforeRequest(scriptPath, []byte(body), req.Model, req.SourceFormat, headers, hostCallbackID)
6468
if errJS != nil {
6569
log.Warnf("failed to execute JS request interceptor [%s]: %v", scriptPath, errJS)
6670
continue
@@ -78,6 +82,10 @@ func (p *jsHandlerPlugin) InterceptRequest(ctx context.Context, req pluginapi.Re
7882
}
7983

8084
func (p *jsHandlerPlugin) InterceptResponse(ctx context.Context, req pluginapi.ResponseInterceptRequest) (pluginapi.ResponseInterceptResponse, error) {
85+
return p.interceptResponse(ctx, req, "")
86+
}
87+
88+
func (p *jsHandlerPlugin) interceptResponse(ctx context.Context, req pluginapi.ResponseInterceptRequest, hostCallbackID string) (pluginapi.ResponseInterceptResponse, error) {
8189
resp := pluginapi.ResponseInterceptResponse{}
8290
scriptPaths := p.allScriptPaths()
8391
if len(scriptPaths) == 0 {
@@ -98,6 +106,7 @@ func (p *jsHandlerPlugin) InterceptResponse(ctx context.Context, req pluginapi.R
98106
scriptPath, req.Model, req.SourceFormat,
99107
reqHeadersMap, req.RequestBody,
100108
bodyStr, nil, respHeaders, false, nil,
109+
hostCallbackID,
101110
)
102111
if errJS != nil {
103112
log.Warnf("failed to execute JS response interceptor [%s]: %v", scriptPath, errJS)
@@ -121,6 +130,10 @@ func (p *jsHandlerPlugin) InterceptResponse(ctx context.Context, req pluginapi.R
121130
}
122131

123132
func (p *jsHandlerPlugin) InterceptStreamChunk(ctx context.Context, req pluginapi.StreamChunkInterceptRequest) (pluginapi.StreamChunkInterceptResponse, error) {
133+
return p.interceptStreamChunk(ctx, req, "")
134+
}
135+
136+
func (p *jsHandlerPlugin) interceptStreamChunk(ctx context.Context, req pluginapi.StreamChunkInterceptRequest, hostCallbackID string) (pluginapi.StreamChunkInterceptResponse, error) {
124137
resp := pluginapi.StreamChunkInterceptResponse{}
125138
scriptPaths := p.allScriptPaths()
126139
if len(scriptPaths) == 0 {
@@ -156,6 +169,7 @@ func (p *jsHandlerPlugin) InterceptStreamChunk(ctx context.Context, req pluginap
156169
scriptPath, req.Model, req.SourceFormat,
157170
reqHeadersMap, req.RequestBody,
158171
"", chunkPtr, respHeaders, !isHeaderInit, historyStrings,
172+
hostCallbackID,
159173
)
160174
if errJS != nil {
161175
log.Warnf("failed to execute JS stream chunk interceptor [%s]: %v", scriptPath, errJS)
@@ -183,13 +197,13 @@ func (p *jsHandlerPlugin) InterceptStreamChunk(ctx context.Context, req pluginap
183197
return resp, nil
184198
}
185199

186-
func (p *jsHandlerPlugin) applyJSBeforeRequest(scriptPath string, payloadBytes []byte, model, protocol string, headers http.Header) ([]byte, []string, error) {
200+
func (p *jsHandlerPlugin) applyJSBeforeRequest(scriptPath string, payloadBytes []byte, model, protocol string, headers http.Header, hostCallbackID string) ([]byte, []string, error) {
187201
program, err := getJSProgram(scriptPath)
188202
if err != nil {
189203
return nil, nil, err
190204
}
191205

192-
engine := newJSEngine()
206+
engine := newJSEngine(newHostJSConsoleLogger(hostCallbackID))
193207
if errRun := engine.runProgram(program, p.cfg.Timeout); errRun != nil {
194208
return nil, nil, errRun
195209
}
@@ -246,13 +260,14 @@ func (p *jsHandlerPlugin) applyJSAfterResponse(
246260
reqHeadersMap map[string]any, reqBody []byte,
247261
bodyStr string, chunkStr *string,
248262
respHeaders http.Header, isStream bool, historyChunks []string,
263+
hostCallbackID string,
249264
) (string, *processedHeaders, bool, error) {
250265
program, err := getJSProgram(scriptPath)
251266
if err != nil {
252267
return bodyStr, nil, false, err
253268
}
254269

255-
engine := newJSEngine()
270+
engine := newJSEngine(newHostJSConsoleLogger(hostCallbackID))
256271
if errRun := engine.runProgram(program, p.cfg.Timeout); errRun != nil {
257272
return bodyStr, nil, false, errRun
258273
}

‎examples/plugin/jshandler/interceptor_test.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ function on_before_request(ctx) {
3131
"gpt-test",
3232
"openai",
3333
headers,
34+
"",
3435
)
3536
if errApply != nil {
3637
t.Fatalf("applyJSBeforeRequest() error = %v", errApply)
@@ -85,6 +86,7 @@ function on_after_stream_response(ctx) {
8586
http.Header{},
8687
true,
8788
[]string{`data: {"choices":[{"delta":{"tool_calls":[{"index":0}]}}]}`},
89+
"",
8890
)
8991
if errApply != nil {
9092
t.Fatalf("applyJSAfterResponse() error = %v", errApply)
@@ -123,6 +125,7 @@ function on_after_nonstream_response(ctx) {
123125
http.Header{},
124126
false,
125127
nil,
128+
"",
126129
)
127130
if errApply != nil {
128131
t.Fatalf("applyJSAfterResponse() error = %v", errApply)

0 commit comments

Comments
 (0)