pluginsdk/plugin/wasmguest/caps/ai.go

package caps

import (
	"context"
	"encoding/json"

	abiv1 "git.dev.alexdunmow.com/block/pluginsdk/abi/v1"
	"git.dev.alexdunmow.com/block/pluginsdk/ai"
)

// aiStub implements ai.ToolRegistry (ai.tools.register) and backs the
// CoreServices.AITextCall func field (ai.text_call).
//
// ToolDefinition.Handler stays guest-side: registration marshals only the
// static descriptor AND records the full definition locally so the host can
// execute the Handler via HOOK_AI_TOOL_CALL. Register tools in Register (every
// pooled instance runs it), not Load (one instance only) — see
// core/docs/wasm-abi.md §"Per-instance state". Instances are single-threaded,
// so the plain map needs no locking.
type aiStub struct {
	base
	tools map[string]*ai.ToolDefinition // slug → definition (with Handler)
}

var _ ai.ToolRegistry = (*aiStub)(nil)

// Register has no error channel; a transport failure is dropped (the host
// records nothing, matching the "best effort at load" contract).
func (s *aiStub) Register(tool *ai.ToolDefinition) {
	if tool == nil {
		return
	}
	if s.tools == nil {
		s.tools = make(map[string]*ai.ToolDefinition)
	}
	s.tools[tool.Slug] = tool
	req := &abiv1.AiToolRegisterRequest{
		Slug:        tool.Slug,
		Name:        tool.Name,
		Description: tool.Description,
	}
	if tool.ParameterSchema != nil {
		if raw, err := json.Marshal(tool.ParameterSchema); err == nil {
			req.ParameterSchemaJson = raw
		}
	}
	_ = s.invoke(context.Background(), "tools.register", req, &abiv1.AiToolRegisterResponse{})
}

// Tool returns the registered tool definition for a slug, for
// HOOK_AI_TOOL_CALL dispatch by the wasm shim.
func (s *aiStub) Tool(slug string) (*ai.ToolDefinition, bool) {
	t, ok := s.tools[slug]
	return t, ok
}

// textCall backs CoreServices.AITextCall.
func (s *aiStub) textCall(ctx context.Context, taskKey, systemPrompt, userMessage string) (string, error) {
	req := &abiv1.AiTextCallRequest{
		TaskKey:      taskKey,
		SystemPrompt: systemPrompt,
		UserMessage:  userMessage,
	}
	resp := &abiv1.AiTextCallResponse{}
	if err := s.invoke(ctx, "text_call", req, resp); err != nil {
		return "", err
	}
	return resp.GetText(), nil
}